From bb6725372b136a362378e53face3db8ab2d53082 Mon Sep 17 00:00:00 2001
From: Mike Oswell <oswell@origamilogic.com>
Date: Fri, 14 Oct 2016 22:36:13 -0700
Subject: [PATCH] Add support for reading CVE data from MySQL.

---
 commands/scan.go     | 15 +++++++++++++--
 config/config.go     | 22 ++++++++++++++++++----
 cveapi/cve_client.go | 11 ++++++++---
 3 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/commands/scan.go b/commands/scan.go
index 1e1df7b746..6c8167f11e 100644
--- a/commands/scan.go
+++ b/commands/scan.go
@@ -45,6 +45,7 @@ type ScanCmd struct {
 	configPath string
 
 	resultsDir       string
+	cvedbtype        string
 	cvedbpath        string
 	cveDictionaryURL string
 	cacheDBPath      string
@@ -90,7 +91,8 @@ func (*ScanCmd) Usage() string {
 		[-lang=en|ja]
 		[-config=/path/to/config.toml]
 		[-results-dir=/path/to/results]
-		[-cve-dictionary-dbpath=/path/to/cve.sqlite3]
+		[-cve-dictionary-dbtype=sqlite3|mysql]
+		[-cve-dictionary-dbpath=/path/to/cve.sqlite3 or mysql connection string]
 		[-cve-dictionary-url=http://127.0.0.1:1323]
 		[-cache-dbpath=/path/to/cache.db]
 		[-cvss-over=7]
@@ -132,6 +134,12 @@ func (p *ScanCmd) SetFlags(f *flag.FlagSet) {
 	defaultResultsDir := filepath.Join(wd, "results")
 	f.StringVar(&p.resultsDir, "results-dir", defaultResultsDir, "/path/to/results")
 
+	f.StringVar(
+		&p.cvedbtype,
+		"cve-dictionary-dbtype",
+		"sqlite3",
+		"DB type for fetching CVE dictionary (sqlite3 or mysql)")
+
 	f.StringVar(
 		&p.cvedbpath,
 		"cve-dictionary-dbpath",
@@ -254,7 +262,9 @@ func (p *ScanCmd) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{})
 	logrus.Info("Start scanning")
 	logrus.Infof("config: %s", p.configPath)
 	if p.cvedbpath != "" {
-		logrus.Infof("cve-dictionary: %s", p.cvedbpath)
+		if p.cvedbtype == "sqlite3" {
+			logrus.Infof("cve-dictionary: %s", p.cvedbpath)
+		}
 	} else {
 		logrus.Infof("cve-dictionary: %s", p.cveDictionaryURL)
 	}
@@ -357,6 +367,7 @@ func (p *ScanCmd) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{})
 	}
 
 	c.Conf.ResultsDir = p.resultsDir
+	c.Conf.CveDBType = p.cvedbtype
 	c.Conf.CveDBPath = p.cvedbpath
 	c.Conf.CveDictionaryURL = p.cveDictionaryURL
 	c.Conf.CacheDBPath = p.cacheDBPath
diff --git a/config/config.go b/config/config.go
index 93c3dffabd..db11a50f8c 100644
--- a/config/config.go
+++ b/config/config.go
@@ -49,6 +49,7 @@ type Config struct {
 
 	HTTPProxy   string `valid:"url"`
 	ResultsDir  string
+	CveDBType   string
 	CveDBPath   string
 	CacheDBPath string
 
@@ -75,10 +76,23 @@ func (c Config) Validate() bool {
 		}
 	}
 
-	if len(c.CveDBPath) != 0 {
-		if ok, _ := valid.IsFilePath(c.CveDBPath); !ok {
-			errs = append(errs, fmt.Errorf(
-				"SQLite3 DB(Cve Dictionary) path must be a *Absolute* file path. -cve-dictionary-dbpath: %s", c.CveDBPath))
+	// If no valid DB type is set, default to sqlite3
+	if c.CveDBType == "" {
+		c.CveDBType = "sqlite3"
+	}
+
+	if c.CveDBType != "sqlite3" && c.CveDBType != "mysql" {
+		errs = append(errs, fmt.Errorf(
+			"CVE DB type must be either 'sqlite3' or 'mysql'.  -cve-dictionary-dbtype: %s", c.CveDBType))
+	}
+
+
+	if c.CveDBType == "sqlite3" {
+		if len(c.CveDBPath) != 0 {
+			if ok, _ := valid.IsFilePath(c.CveDBPath); !ok {
+				errs = append(errs, fmt.Errorf(
+					"SQLite3 DB(Cve Dictionary) path must be a *Absolute* file path. -cve-dictionary-dbpath: %s", c.CveDBPath))
+			}
 		}
 	}
 
diff --git a/cveapi/cve_client.go b/cveapi/cve_client.go
index eb33838183..44f5f282b6 100644
--- a/cveapi/cve_client.go
+++ b/cveapi/cve_client.go
@@ -49,7 +49,7 @@ func (api *cvedictClient) initialize() {
 
 func (api cvedictClient) CheckHealth() (ok bool, err error) {
 	if config.Conf.CveDBPath != "" {
-		log.Debugf("get cve-dictionary from sqlite3")
+		log.Debugf("get cve-dictionary from %s", config.Conf.CveDBType)
 		return true, nil
 	}
 
@@ -135,8 +135,10 @@ func (api cvedictClient) FetchCveDetails(cveIDs []string) (cveDetails cve.CveDet
 }
 
 func (api cvedictClient) FetchCveDetailsFromCveDB(cveIDs []string) (cveDetails cve.CveDetails, err error) {
-	log.Debugf("open cve-dictionary db")
+	log.Debugf("open cve-dictionary db (%s)", config.Conf.CveDBType)
+	cveconfig.Conf.DBType = config.Conf.CveDBType
 	cveconfig.Conf.DBPath = config.Conf.CveDBPath
+	cveconfig.Conf.DebugSQL = config.Conf.DebugSQL
 	if err := cvedb.OpenDB(); err != nil {
 		return []cve.CveDetail{},
 			fmt.Errorf("Failed to open DB. err: %s", err)
@@ -239,8 +241,11 @@ func (api cvedictClient) httpPost(key, url string, query map[string]string) ([]c
 }
 
 func (api cvedictClient) FetchCveDetailsByCpeNameFromDB(cpeName string) ([]cve.CveDetail, error) {
-	log.Debugf("open cve-dictionary db")
+	log.Debugf("open cve-dictionary db (%s)", config.Conf.CveDBType)
+	cveconfig.Conf.DBType = config.Conf.CveDBType
 	cveconfig.Conf.DBPath = config.Conf.CveDBPath
+	cveconfig.Conf.DebugSQL = config.Conf.DebugSQL
+
 	if err := cvedb.OpenDB(); err != nil {
 		return []cve.CveDetail{},
 			fmt.Errorf("Failed to open DB. err: %s", err)