dbstore/dbstore.go: add AddCisaCatalogs() and CisaSearch()

1 files changed, 154 insertions, 0 deletions

diff --git a/dbstore/dbstore.go b/dbstore/dbstore.go
index 7132658..08da511 100644
--- a/dbstore/dbstore.go
+++ b/dbstore/dbstore.go
@@ -7,6 +7,7 @@ import (
   "encoding/json"
   "fmt"
   _ "github.com/mattn/go-sqlite3"
+  "github.com/pablotron/cvez/cisa"
   "github.com/pablotron/cvez/cpedict"
   "github.com/pablotron/cvez/cpematch"
   nvd_feed "github.com/pablotron/cvez/feed"
@@ -750,3 +751,156 @@ func (me DbStore) CveSearch(
   // return results
   return r, err
 }
+
+// Add CISA vendor.
+//
+// Note: This query has an ON CONFLICT(name) DO NOTHING clause, so if
+// the name already exists it does nothing.
+func (me DbStore) addCisaVendor(ctx context.Context, tx Tx, name string) error {
+  _, err := tx.Exec(ctx, "cisa/insert-vendor", string(name))
+  return err
+}
+
+// Add CISA product.
+//
+// Note: This query has an ON CONFLICT(name) DO NOTHING clause, so if
+// the name already exists it does nothing.
+func (me DbStore) addCisaProduct(ctx context.Context, tx Tx, name string) error {
+  _, err := tx.Exec(ctx, "cisa/insert-product", string(name))
+  return err
+}
+
+// Add CISA vulnerability and return vulnerability ID
+func (me DbStore) addCisaVuln(ctx context.Context, tx Tx, catId int64, vuln cisa.Vulnerability) error {
+  // add vendor (if it does not already exist)
+  if err := me.addCisaVendor(ctx, tx, vuln.VendorProject); err != nil {
+    return err
+  }
+
+  // add product (if it does not already exist)
+  if err := me.addCisaProduct(ctx, tx, vuln.Product); err != nil {
+    return err
+  }
+
+  // exec query
+  _, err := tx.Exec(ctx, "cisa/insert-vuln",
+    catId,
+    vuln.CveId.Year(),
+    vuln.CveId.Number(),
+    vuln.VendorProject,
+    vuln.Product,
+    vuln.Name,
+    vuln.DateAdded.String(),
+    vuln.ShortDescription,
+    vuln.RequiredAction,
+    vuln.DueDate.String(),
+  )
+
+  // return result
+  return err
+}
+
+// Add CISA catalog and return catalog ID
+func (me DbStore) addCisaCatalog(ctx context.Context, tx Tx, cat cisa.Catalog) (int64, error) {
+  // return catalog ID
+  var catId int64
+
+  // serialize date released
+  ts, err := cat.DateReleased.MarshalText()
+  if err != nil {
+    return catId, err
+  }
+
+  // add catalog
+  rs, err := tx.Exec(ctx, "cisa/insert", cat.Title, cat.Version, string(ts))
+  if err != nil {
+    return catId, err
+  }
+
+  // get catalog ID
+  catId, err = rs.LastInsertId()
+  if err != nil {
+    return catId, err
+  }
+
+  // add catalog vulnerabilities
+  for _, vuln := range(cat.Vulnerabilities) {
+    if err := me.addCisaVuln(ctx, tx, catId, vuln); err != nil {
+      return catId, err
+    }
+  }
+
+  // return success
+  return catId, nil
+}
+
+// query IDs used by AddCisaCatalogs()
+var addCisaCatalogsQueryIds = []string {
+  "cisa/insert",
+  "cisa/insert-vendor",
+  "cisa/insert-product",
+  "cisa/insert-vuln",
+  "cisa/insert-fts-refresh",
+}
+
+// Refresh CVE FTS index.
+func (me DbStore) refreshCisaFts(ctx context.Context, tx Tx) error {
+  _, err := tx.Exec(ctx, "cisa/insert-fts-refresh")
+  return err
+}
+
+// Import CISA Known Exploited Vulnerabilities catalogs.
+func (me DbStore) AddCisaCatalogs(ctx context.Context, cats []cisa.Catalog) ([]int64, error) {
+  ids := make([]int64, len(cats))
+
+  // lazy-init db
+  if err := me.Init(ctx); err != nil {
+    return ids, err
+  }
+
+  // begin transaction
+  err := me.Tx(ctx, addCisaCatalogsQueryIds, func(tx Tx) error {
+    for i, cat := range(cats) {
+      // add feed, get feed ID
+      if id, err := me.addCisaCatalog(ctx, tx, cat); err != nil {
+        return err
+      } else {
+        ids[i] = id
+      }
+    }
+
+    // refresh the cve fts index
+    return me.refreshCisaFts(ctx, tx)
+  })
+
+  // return results
+  return ids, err
+}
+
+// search CISA vulnerabilities.
+func (me DbStore) CisaSearch(
+  ctx context.Context,
+  s string,
+) ([]CisaSearchRow, error) {
+  var r []CisaSearchRow
+
+  // lazy-init db
+  if err := me.Init(ctx); err != nil {
+    return r, err
+  }
+
+  // get/exec search query
+  err := me.QueryRow(ctx, "cisa/search", []interface{} {
+    db_sql.Named("q", s),
+  }, func(row *db_sql.Row) error {
+    var s string
+    if err := row.Scan(&s); err != nil {
+      return err
+    }
+
+    return json.Unmarshal([]byte(s), &r)
+  })
+
+  // return results
+  return r, err
+}