feat(fetch): add vulncheck nvd++

Author: MaineK00n

commands/fetchvulncheck.go

@@ -0,0 +1,70 @@
+package commands
+
+import (
+	"errors"
+	"time"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+	"golang.org/x/xerrors"
+
+	"github.com/vulsio/go-cve-dictionary/db"
+	"github.com/vulsio/go-cve-dictionary/fetcher/vulncheck"
+	"github.com/vulsio/go-cve-dictionary/log"
+	"github.com/vulsio/go-cve-dictionary/models"
+)
+
+var fetchVulncheckCmd = &cobra.Command{
+	Use:   "vulncheck",
+	Short: "Fetch Vulnerability dictionary from VulnCheck NVD++",
+	Long:  "Fetch Vulnerability dictionary from VulnCheck NVD++",
+	RunE:  fetchVulncheck,
+}
+
+func init() {
+	fetchCmd.AddCommand(fetchVulncheckCmd)
+}
+
+func fetchVulncheck(_ *cobra.Command, args []string) (err error) {
+	if err := log.SetLogger(viper.GetBool("log-to-file"), viper.GetString("log-dir"), viper.GetBool("debug"), viper.GetBool("log-json")); err != nil {
+		return xerrors.Errorf("Failed to SetLogger. err: %w", err)
+	}
+
+	driver, err := db.NewDB(viper.GetString("dbtype"), viper.GetString("dbpath"), viper.GetBool("debug-sql"), db.Option{})
+	if err != nil {
+		if errors.Is(err, db.ErrDBLocked) {
+			return xerrors.Errorf("Failed to open DB. Close DB connection before fetching. err: %w", err)
+		}
+		return xerrors.Errorf("Failed to open DB. err: %w", err)
+	}
+
+	fetchMeta, err := driver.GetFetchMeta()
+	if err != nil {
+		return xerrors.Errorf("Failed to get FetchMeta from DB. err: %w", err)
+	}
+	if fetchMeta.OutDated() {
+		return xerrors.Errorf("Failed to Insert CVEs into DB. err: SchemaVersion is old. SchemaVersion: %+v", map[string]uint{"latest": models.LatestSchemaVersion, "DB": fetchMeta.SchemaVersion})
+	}
+	// If the fetch fails the first time (without SchemaVersion), the DB needs to be cleaned every time, so insert SchemaVersion.
+	if err := driver.UpsertFetchMeta(fetchMeta); err != nil {
+		return xerrors.Errorf("Failed to upsert FetchMeta to DB. dbpath: %s, err: %w", viper.GetString("dbpath"), err)
+	}
+
+	cves, err := vulncheck.FetchConvert()
+	if err != nil {
+		return xerrors.Errorf("Failed to fetch from VulnCheck NVD++. err: %w", err)
+	}
+
+	log.Infof("Inserting VulnCheck NVD++ into DB (%s).", driver.Name())
+	if err := driver.InsertVulncheck(cves); err != nil {
+		return xerrors.Errorf("Failed to insert. dbpath: %s, err: %w", viper.GetString("dbpath"), err)
+	}
+
+	fetchMeta.LastFetchedAt = time.Now()
+	if err := driver.UpsertFetchMeta(fetchMeta); err != nil {
+		return xerrors.Errorf("Failed to upsert FetchMeta to DB. dbpath: %s, err: %w", viper.GetString("dbpath"), err)
+	}
+
+	log.Infof("Finished fetching VulnCheck NVD++.")
+	return nil
+}

db/db.go

@@ -41,12 +41,14 @@ type DB interface {
 	GetAdvisoriesCisco() (map[string][]string, error)
 	InsertJvn([]string) error
 	InsertNvd([]string) error
+	InsertVulncheck(iter.Seq2[models.Vulncheck, error]) error
 	InsertFortinet([]models.Fortinet) error
 	InsertMitre([]string) error
 	InsertPaloalto(iter.Seq2[models.Paloalto, error]) error
 	InsertCisco(iter.Seq2[models.Cisco, error]) error
-	CountNvd() (int, error)
 	CountJvn() (int, error)
+	CountNvd() (int, error)
+	CountVulncheck() (int, error)
 	CountFortinet() (int, error)
 	CountMitre() (int, error)
 	CountPaloalto() (int, error)

db/rdb.go

@@ -160,6 +160,16 @@ func (r *RDBDriver) MigrateDB() error {
 		&models.NvdReference{},
 		&models.NvdCert{},
 
+		&models.Vulncheck{},
+		&models.VulncheckDescription{},
+		&models.VulncheckCvss2Extra{},
+		&models.VulncheckCvss3{},
+		&models.VulncheckCvss40{},
+		&models.VulncheckCwe{},
+		&models.VulncheckCpe{},
+		&models.VulncheckReference{},
+		&models.VulncheckCert{},
+
 		&models.Jvn{},
 		&models.JvnCvss2{},
 		&models.JvnCvss3{},
@@ -894,6 +904,85 @@ func insertNvd(tx *gorm.DB, cves []models.Nvd, batchSize int) error {
 	return nil
 }
 
+// CountVulncheck count vulncheck table
+func (r *RDBDriver) CountVulncheck() (int, error) {
+	var count int64
+	if err := r.conn.Model(&models.Vulncheck{}).Count(&count).Error; err != nil {
+		return 0, err
+	}
+	return int(count), nil
+}
+
+// InsertVulncheck Cve information from DB.
+func (r *RDBDriver) InsertVulncheck(cves iter.Seq2[models.Vulncheck, error]) (err error) {
+	bar := pb.ProgressBarTemplate(`{{cycle . "[                    ]" "[=>                  ]" "[===>                ]" "[=====>              ]" "[======>             ]" "[========>           ]" "[==========>         ]" "[============>       ]" "[==============>     ]" "[================>   ]" "[==================> ]" "[===================>]"}} {{counters .}} files processed. ({{speed .}})`).New(0).Start().SetWriter(func() io.Writer {
+		if viper.GetBool("log-json") {
+			return io.Discard
+		}
+		return os.Stderr
+	}())
+
+	tx := r.conn.Begin()
+	defer func() {
+		if re := recover(); re != nil {
+			tx.Rollback()
+		}
+	}()
+	if err := tx.Error; err != nil {
+		return err
+	}
+
+	batchSize := viper.GetInt("batch-size")
+	if batchSize < 1 {
+		return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly")
+	}
+
+	logger.Infof("Deleting Vulncheck tables...")
+	if err := deleteVulncheck(tx); err != nil {
+		tx.Rollback()
+		return xerrors.Errorf("Failed to deleteVulncheck. err: %w", err)
+	}
+
+	for cve, err := range cves {
+		if err != nil {
+			tx.Rollback()
+			return xerrors.Errorf("Failed to convert Vulncheck cves. err: %w", err)
+		}
+
+		if err := tx.Omit("Cpes").Create(&cve).Error; err != nil {
+			tx.Rollback()
+			return xerrors.Errorf("Failed to insert. err: %w", err)
+		}
+
+		for i := range cve.Cpes {
+			cve.Cpes[i].VulncheckID = uint(cve.ID)
+		}
+
+		for chunk := range slices.Chunk(cve.Cpes, batchSize) {
+			if err := tx.Create(chunk).Error; err != nil {
+				tx.Rollback()
+				return xerrors.Errorf("Failed to insert. err: %w", err)
+			}
+		}
+
+		bar.Increment()
+	}
+
+	if err := tx.Commit().Error; err != nil {
+		return xerrors.Errorf("Failed to Commit Transaction. err: %w", err)
+	}
+	return nil
+}
+
+func deleteVulncheck(tx *gorm.DB) error {
+	for _, table := range []interface{}{models.Vulncheck{}, models.VulncheckDescription{}, models.VulncheckCvss2Extra{}, models.VulncheckCvss3{}, models.VulncheckCvss40{}, models.VulncheckCwe{}, models.VulncheckCpe{}, models.VulncheckReference{}, models.VulncheckCert{}} {
+		if err := tx.Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(table).Error; err != nil {
+			return xerrors.Errorf("Failed to delete old records. err: %w", err)
+		}
+	}
+	return nil
+}
+
 // GetAdvisoriesFortinet get AdvisoryID: []CVE IDs
 func (r *RDBDriver) GetAdvisoriesFortinet() (map[string][]string, error) {
 	var fs []models.Fortinet