github.com/khulnasoft-lab/tunnel-db@v0.0.0-20231117205118-74e1113bd007/pkg/vulnsrc/nvd/nvd.go (about)

     1  package nvd
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"io"
     7  	"log"
     8  	"path/filepath"
     9  	"strings"
    10  	"time"
    11  
    12  	bolt "go.etcd.io/bbolt"
    13  	"golang.org/x/xerrors"
    14  
    15  	"github.com/khulnasoft-lab/tunnel-db/pkg/db"
    16  	"github.com/khulnasoft-lab/tunnel-db/pkg/types"
    17  	"github.com/khulnasoft-lab/tunnel-db/pkg/utils"
    18  	"github.com/khulnasoft-lab/tunnel-db/pkg/vulnsrc/vulnerability"
    19  )
    20  
    21  const (
    22  	vulnListDir = "vuln-list-nvd"
    23  	feedDir     = "feed"
    24  )
    25  
    26  type VulnSrc struct {
    27  	dbc db.Operation
    28  }
    29  
    30  func NewVulnSrc() VulnSrc {
    31  	return VulnSrc{
    32  		dbc: db.Config{},
    33  	}
    34  }
    35  
    36  func (vs VulnSrc) Name() types.SourceID {
    37  	return vulnerability.NVD
    38  }
    39  
    40  func (vs VulnSrc) Update(dir string) error {
    41  	rootDir := filepath.Join(dir, vulnListDir, feedDir)
    42  
    43  	var items []Item
    44  	buffer := &bytes.Buffer{}
    45  	err := utils.FileWalk(rootDir, func(r io.Reader, _ string) error {
    46  		item := Item{}
    47  		if _, err := buffer.ReadFrom(r); err != nil {
    48  			return xerrors.Errorf("failed to read file: %w", err)
    49  		}
    50  		if err := json.Unmarshal(buffer.Bytes(), &item); err != nil {
    51  			return xerrors.Errorf("failed to decode NVD JSON: %w", err)
    52  		}
    53  		buffer.Reset()
    54  		items = append(items, item)
    55  		return nil
    56  	})
    57  	if err != nil {
    58  		return xerrors.Errorf("error in NVD walk: %w", err)
    59  	}
    60  
    61  	if err = vs.save(items); err != nil {
    62  		return xerrors.Errorf("error in NVD save: %w", err)
    63  	}
    64  
    65  	return nil
    66  }
    67  
    68  func (vs VulnSrc) commit(tx *bolt.Tx, items []Item) error {
    69  	for _, item := range items {
    70  		cveID := item.Cve.Meta.ID
    71  		severity, _ := types.NewSeverity(item.Impact.BaseMetricV2.Severity)
    72  		severityV3, _ := types.NewSeverity(item.Impact.BaseMetricV3.CvssV3.BaseSeverity)
    73  
    74  		var references []string
    75  		for _, ref := range item.Cve.References.ReferenceDataList {
    76  			references = append(references, ref.URL)
    77  		}
    78  
    79  		var (
    80  			description string
    81  		)
    82  		for _, d := range item.Cve.Description.DescriptionDataList {
    83  			if d.Value != "" {
    84  				description = d.Value
    85  				break
    86  			}
    87  		}
    88  		var cweIDs []string
    89  		for _, data := range item.Cve.ProblemType.ProblemTypeData {
    90  			for _, desc := range data.Description {
    91  				if !strings.HasPrefix(desc.Value, "CWE") {
    92  					continue
    93  				}
    94  				cweIDs = append(cweIDs, desc.Value)
    95  			}
    96  		}
    97  
    98  		publishedDate, _ := time.Parse("2006-01-02T15:04Z", item.PublishedDate)
    99  		lastModifiedDate, _ := time.Parse("2006-01-02T15:04Z", item.LastModifiedDate)
   100  
   101  		vuln := types.VulnerabilityDetail{
   102  			CvssScore:        item.Impact.BaseMetricV2.CvssV2.BaseScore,
   103  			CvssVector:       item.Impact.BaseMetricV2.CvssV2.VectorString,
   104  			CvssScoreV3:      item.Impact.BaseMetricV3.CvssV3.BaseScore,
   105  			CvssVectorV3:     item.Impact.BaseMetricV3.CvssV3.VectorString,
   106  			Severity:         severity,
   107  			SeverityV3:       severityV3,
   108  			CweIDs:           cweIDs,
   109  			References:       references,
   110  			Title:            "",
   111  			Description:      description,
   112  			PublishedDate:    &publishedDate,
   113  			LastModifiedDate: &lastModifiedDate,
   114  		}
   115  
   116  		if err := vs.dbc.PutVulnerabilityDetail(tx, cveID, vulnerability.NVD, vuln); err != nil {
   117  			return err
   118  		}
   119  	}
   120  	return nil
   121  }
   122  
   123  func (vs VulnSrc) save(items []Item) error {
   124  	log.Println("NVD batch update")
   125  	err := vs.dbc.BatchUpdate(func(tx *bolt.Tx) error {
   126  		return vs.commit(tx, items)
   127  	})
   128  	if err != nil {
   129  		return xerrors.Errorf("error in batch update: %w", err)
   130  	}
   131  	return nil
   132  }