dbstore/dbstore.go: add Update()

1 files changed, 259 insertions, 107 deletions

diff --git a/dbstore/dbstore.go b/dbstore/dbstore.go
index 65aa359..fca85fc 100644
--- a/dbstore/dbstore.go
+++ b/dbstore/dbstore.go
@@ -10,7 +10,11 @@ import (
   "github.com/pablotron/cvez/cisa"
   "github.com/pablotron/cvez/cpedict"
   "github.com/pablotron/cvez/cpematch"
+  "github.com/pablotron/cvez/nvdmirror"
+  "github.com/pablotron/cvez/util"
   nvd_feed "github.com/pablotron/cvez/feed"
+  "github.com/rs/zerolog/log"
+  "path/filepath"
   "sort"
 )
 
@@ -185,47 +189,52 @@ func (me DbStore) refreshCpeFts(ctx context.Context, tx Tx) error {
   return err
 }
 
-// import CPE dictionary
-func (me DbStore) AddCpeDictionary(ctx context.Context, dict cpedict.Dictionary) error {
-  // lazy-init db
-  if err := me.Init(ctx); err != nil {
-    return err
-  }
+// import CPE dictionary within transaction.
+func (me DbStore) addCpeDictionary(ctx context.Context, tx Tx, dict cpedict.Dictionary) error {
+  // add items
+  for _, item := range(dict.Items) {
+    // add cpe
+    rs, err := tx.Exec(ctx, "cpe/insert", item.CpeUri, item.Cpe23Item.Name)
+    if err != nil {
+      return err
+    }
 
-  return me.Tx(ctx, addCpeDictionaryQueryIds, func(tx Tx) error {
-    // add items
-    for _, item := range(dict.Items) {
-      // add cpe
-      rs, err := tx.Exec(ctx, "cpe/insert", item.CpeUri, item.Cpe23Item.Name)
+    // get last row ID
+    id, err := rs.LastInsertId()
+    if err != nil {
+      return err
+    }
+
+    // add titles
+    for _, title := range(item.Titles) {
+      _, err := tx.Exec(ctx, "cpe/insert-title", id, title.Lang, title.Text)
       if err != nil {
         return err
       }
+    }
 
-      // get last row ID
-      id, err := rs.LastInsertId()
+    // add refs
+    for _, ref := range(item.References) {
+      _, err := tx.Exec(ctx, "cpe/insert-ref", id, ref.Href, ref.Text)
       if err != nil {
         return err
       }
+    }
+  }
 
-      // add titles
-      for _, title := range(item.Titles) {
-        _, err := tx.Exec(ctx, "cpe/insert-title", id, title.Lang, title.Text)
-        if err != nil {
-          return err
-        }
-      }
+  // refresh the cpe fts, return result
+  return me.refreshCpeFts(ctx, tx)
+}
 
-      // add refs
-      for _, ref := range(item.References) {
-        _, err := tx.Exec(ctx, "cpe/insert-ref", id, ref.Href, ref.Text)
-        if err != nil {
-          return err
-        }
-      }
-    }
+// import CPE dictionary
+func (me DbStore) AddCpeDictionary(ctx context.Context, dict cpedict.Dictionary) error {
+  // lazy-init db
+  if err := me.Init(ctx); err != nil {
+    return err
+  }
 
-    // refresh the cpe fts's
-    return me.refreshCpeFts(ctx, tx)
+  return me.Tx(ctx, addCpeDictionaryQueryIds, func(tx Tx) error {
+    return me.addCpeDictionary(ctx, tx, dict)
   })
 }
 
@@ -269,78 +278,83 @@ var addCpeMatchesQueryIds = []string {
   "cpe-match/insert-name",
 }
 
-// import CPE matches
-func (me DbStore) AddCpeMatches(ctx context.Context, matches cpematch.Matches) error {
-  // lazy-init db
-  if err := me.Init(ctx); err != nil {
-    return err
-  }
+// import CPE matches within transaction
+func (me DbStore) addCpeMatches(ctx context.Context, tx Tx, matches cpematch.Matches) error {
+  // add matches
+  for _, m := range(matches.Matches) {
+    // add cpe
+    rs, err := tx.Exec(ctx, "cpe-match/insert", m.Cpe23Uri, m.Cpe22Uri)
+    if err != nil {
+      return err
+    }
 
-  // begin transaction
-  return me.Tx(ctx, addCpeMatchesQueryIds, func(tx Tx) error {
-    // add matches
-    for _, m := range(matches.Matches) {
-      // add cpe
-      rs, err := tx.Exec(ctx, "cpe-match/insert", m.Cpe23Uri, m.Cpe22Uri)
-      if err != nil {
+    // get last row ID
+    id, err := rs.LastInsertId()
+    if err != nil {
+      return err
+    }
+
+    // add vulnerable
+    if m.Vulnerable != nil {
+      _, err := tx.Exec(ctx, "cpe-match/insert-vulnerable", id, *m.Vulnerable)
+      if  err != nil {
         return err
       }
+    }
 
-      // get last row ID
-      id, err := rs.LastInsertId()
-      if err != nil {
+    // add version minimum
+    if m.VersionStartIncluding != "" && m.VersionStartExcluding != "" {
+      return fmt.Errorf("cannot specify both VersionStartIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionStartIncluding, m.VersionStartExcluding)
+    } else if m.VersionStartIncluding != "" {
+      _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, true, m.VersionStartIncluding)
+      if  err != nil {
         return err
       }
-
-      // add vulnerable
-      if m.Vulnerable != nil {
-        _, err := tx.Exec(ctx, "cpe-match/insert-vulnerable", id, *m.Vulnerable)
-        if  err != nil {
-          return err
-        }
+    } else if m.VersionStartExcluding != "" {
+      _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, false, m.VersionStartExcluding)
+      if  err != nil {
+        return err
       }
+    }
 
-      // add version minimum
-      if m.VersionStartIncluding != "" && m.VersionStartExcluding != "" {
-        return fmt.Errorf("cannot specify both VersionStartIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionStartIncluding, m.VersionStartExcluding)
-      } else if m.VersionStartIncluding != "" {
-        _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, true, m.VersionStartIncluding)
-        if  err != nil {
-          return err
-        }
-      } else if m.VersionStartExcluding != "" {
-        _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, false, m.VersionStartExcluding)
-        if  err != nil {
-          return err
-        }
+    // add version maximum
+    if m.VersionEndIncluding != "" && m.VersionEndExcluding != "" {
+      return fmt.Errorf("cannot specify both VersionEndIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionEndIncluding, m.VersionEndExcluding)
+    } else if m.VersionEndIncluding != "" {
+      _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, true, m.VersionEndIncluding)
+      if  err != nil {
+        return err
       }
-
-      // add version maximum
-      if m.VersionEndIncluding != "" && m.VersionEndExcluding != "" {
-        return fmt.Errorf("cannot specify both VersionEndIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionEndIncluding, m.VersionEndExcluding)
-      } else if m.VersionEndIncluding != "" {
-        _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, true, m.VersionEndIncluding)
-        if  err != nil {
-          return err
-        }
-      } else if m.VersionEndExcluding != "" {
-        _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, false, m.VersionEndExcluding)
-        if  err != nil {
-          return err
-        }
+    } else if m.VersionEndExcluding != "" {
+      _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, false, m.VersionEndExcluding)
+      if  err != nil {
+        return err
       }
+    }
 
-      // add names
-      for _, name := range(m.Names) {
-        _, err := tx.Exec(ctx, "cpe-match/insert-name", id, name.Cpe22Uri, name.Cpe23Uri)
-        if err != nil {
-          return err
-        }
+    // add names
+    for _, name := range(m.Names) {
+      _, err := tx.Exec(ctx, "cpe-match/insert-name", id, name.Cpe22Uri, name.Cpe23Uri)
+      if err != nil {
+        return err
       }
     }
+  }
 
-    // return success
-    return nil
+  // return success
+  return nil
+}
+
+// import CPE matches
+func (me DbStore) AddCpeMatches(ctx context.Context, matches cpematch.Matches) error {
+  // lazy-init db
+  if err := me.Init(ctx); err != nil {
+    return err
+  }
+
+  // begin transaction
+  return me.Tx(ctx, addCpeMatchesQueryIds, func(tx Tx) error {
+    return me.addCpeMatches(ctx, tx, matches)
   })
 }
 
@@ -695,6 +709,23 @@ func (me DbStore) refreshCveFts(ctx context.Context, tx Tx) error {
   return err
 }
 
+// Import CVE feeds within transaction.
+func (me DbStore) addCveFeeds(ctx context.Context, tx Tx, feeds []nvd_feed.Feed) ([]int64, error) {
+  feedIds := make([]int64, len(feeds))
+
+  for i, feed := range(feeds) {
+    // add feed, get feed ID
+    if id, err := me.addFeed(ctx, tx, feed); err != nil {
+      return feedIds, err
+    } else {
+      feedIds[i] = id
+    }
+  }
+
+  // refresh the cve fts index, return results
+  return feedIds, me.refreshCveFts(ctx, tx)
+}
+
 // Import CVE feeds.
 func (me DbStore) AddCveFeeds(ctx context.Context, feeds []nvd_feed.Feed) ([]int64, error) {
   feedIds := make([]int64, len(feeds))
@@ -706,17 +737,16 @@ func (me DbStore) AddCveFeeds(ctx context.Context, feeds []nvd_feed.Feed) ([]int
 
   // begin transaction
   err := me.Tx(ctx, addCveFeedsQueryIds, func(tx Tx) error {
-    for i, feed := range(feeds) {
-      // add feed, get feed ID
-      if id, err := me.addFeed(ctx, tx, feed); err != nil {
-        return err
-      } else {
+    if ids, err := me.addCveFeeds(ctx, tx, feeds); err != nil {
+      return err
+    } else {
+      for i, id := range(ids) {
         feedIds[i] = id
       }
-    }
 
-    // refresh the cve fts index
-    return me.refreshCveFts(ctx, tx)
+      // return success
+      return nil
+    }
   })
 
   // return results
@@ -857,6 +887,23 @@ func (me DbStore) refreshCisaFts(ctx context.Context, tx Tx) error {
   return err
 }
 
+// Import CISA catalog within transaction
+func (me DbStore) addCisaCatalogs(ctx context.Context, tx Tx, cats []cisa.Catalog) ([]int64, error) {
+  ids := make([]int64, len(cats))
+
+  for i, cat := range(cats) {
+    // add catalog, get catalog ID
+    if id, err := me.addCisaCatalog(ctx, tx, cat); err != nil {
+      return ids, err
+    } else {
+      ids[i] = id
+    }
+  }
+
+  // refresh the cisa fts index
+  return ids, me.refreshCisaFts(ctx, tx)
+}
+
 // Import CISA Known Exploited Vulnerabilities catalogs.
 func (me DbStore) AddCisaCatalogs(ctx context.Context, cats []cisa.Catalog) ([]int64, error) {
   ids := make([]int64, len(cats))
@@ -868,17 +915,8 @@ func (me DbStore) AddCisaCatalogs(ctx context.Context, cats []cisa.Catalog) ([]i
 
   // 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)
+    _, err := me.addCisaCatalogs(ctx, tx, cats)
+    return err
   })
 
   // return results
@@ -915,3 +953,117 @@ func (me DbStore) CisaSearch(
   // return results
   return r, err
 }
+
+var updateQueryIds = []string {
+  // addCveFeeds
+  "feed/insert",
+  "feed/insert-assigner",
+  "feed/insert-cve",
+  "feed/insert-cve-desc",
+  "feed/insert-cve-fts-refresh",
+  "feed/insert-cve-problem",
+  "feed/insert-cve-problem-desc",
+  "feed/insert-cve-ref",
+  "feed/insert-cve-ref-tag",
+  "feed/insert-desc",
+  "feed/insert-item",
+  "feed/insert-item-cvss-v2",
+  "feed/insert-item-cvss-v3",
+
+  // addCpeDictionary
+  "cpe/insert",
+  "cpe/insert-title",
+  "cpe/insert-ref",
+  "cpe/insert-fts-refresh",
+
+  // addCpeMatchesQueryIds
+  "cpe-match/insert",
+  "cpe-match/insert-vulnerable",
+  "cpe-match/insert-version-min",
+  "cpe-match/insert-version-max",
+  "cpe-match/insert-name",
+
+  // addCisaCatalogQueryIds
+  "cisa/insert",
+  "cisa/insert-vendor",
+  "cisa/insert-product",
+  "cisa/insert-vuln",
+  "cisa/insert-fts-refresh",
+}
+
+// Apply updates from cache directory.
+func (me DbStore) Update(ctx context.Context, cacheDir string, updates []nvdmirror.Update) error {
+  // lazy-init db
+  if err := me.Init(ctx); err != nil {
+    return err
+  }
+
+  return me.Tx(ctx, updateQueryIds, func(tx Tx) error {
+    // load feeds
+    log.Info().Msg("getFeeds")
+    if feeds, err := util.GetFeeds(cacheDir, updates); err != nil {
+      return err
+    } else if len(feeds) > 0 {
+      // add feeds
+      log.Info().Msg("addCveFeeds")
+      if _, err = me.addCveFeeds(ctx, tx, feeds); err != nil {
+        return err
+      }
+    }
+
+    // process cpe dictionary before cpe matches to prevent FK
+    // constraint violations
+    for _, row := range(updates) {
+      if row.Type == nvdmirror.UpdateCpeDict {
+        // load dictionary
+        log.Info().Msg("getCpeDict")
+        if dict, err := util.GetCpeDict(filepath.Join(cacheDir, row.Path)); err != nil {
+          return err
+        } else {
+          // add dict
+          log.Info().Msg("addCpeDictionary")
+          if err = me.addCpeDictionary(ctx, tx, dict); err != nil {
+            return err
+          }
+        }
+      }
+    }
+
+    // process remaining updates
+    for _, row := range(updates) {
+      switch row.Type {
+      case nvdmirror.UpdateCpeMatch:
+        log.Info().Msg("getCpeMatches")
+        matches, err := util.GetCpeMatches(filepath.Join(cacheDir, row.Path))
+        if err != nil {
+          return err
+        }
+
+
+        log.Info().Msg("addCpeMatches")
+        if err = me.addCpeMatches(ctx, tx, matches); err != nil {
+          return err
+        }
+      case nvdmirror.UpdateCisaKevc:
+        log.Info().Msg("getCisaCatalog")
+        cat, err := util.GetCisaCatalog(filepath.Join(cacheDir, row.Path))
+        if err != nil {
+          return err
+        }
+
+        log.Info().Msg("addCisaCatalog")
+        if _, err = me.addCisaCatalog(ctx, tx, cat); err != nil {
+          return err
+        }
+
+        // refresh the cisa fts index
+        if err = me.refreshCisaFts(ctx, tx); err != nil {
+          return err
+        }
+      }
+    }
+
+    // return success
+    return nil
+  })
+}