dbstore/dbstore.go: add AddCveFeed() (incomplete)

1 files changed, 256 insertions, 0 deletions

diff --git a/dbstore/dbstore.go b/dbstore/dbstore.go
index ccddd72..b6ba122 100644
--- a/dbstore/dbstore.go
+++ b/dbstore/dbstore.go
@@ -8,6 +8,7 @@ import (
   _ "github.com/mattn/go-sqlite3"
   "github.com/pablotron/cvez/cpedict"
   "github.com/pablotron/cvez/cpematch"
+  nvd_feed "github.com/pablotron/cvez/feed"
 )
 
 // sqlite3 backing store
@@ -345,3 +346,258 @@ func (me DbStore) CpeMatchSearch(
   // return result
   return r, err
 }
+
+// Check data type, data format, and data version.
+func (me DbStore) checkData(
+  dataType nvd_feed.DataType,
+  dataFormat nvd_feed.DataFormat,
+  dataVersion nvd_feed.DataVersion,
+) error {
+  // check data type
+  if dataType != nvd_feed.CveType {
+    return fmt.Errorf("unknown data type: %s", dataType)
+  }
+
+  // check data format
+  if dataFormat != nvd_feed.MitreFormat {
+    return fmt.Errorf("unknown data format: %s", dataFormat)
+  }
+
+  // check data version
+  if dataVersion != nvd_feed.V40 {
+    return fmt.Errorf("unknown data version: %s", dataVersion)
+  }
+
+  // return success
+  return nil
+}
+
+// Add description.
+func (me DbStore) addDescriptions(ctx context.Context, tx Tx, ds []nvd_feed.Description) ([]int64, error) {
+  r := make([]int64, len(ds))
+
+  for i, d := range(ds) {
+    // add description
+    rs, err := tx.Exec(ctx, "feed/insert-description", d.Lang, d.Value)
+    if err != nil {
+      return r, err
+    }
+
+    // return id and error
+    id, err := rs.LastInsertId()
+    if err != nil {
+      return r, err
+    }
+
+    // add to results
+    r[i] = id
+  }
+
+  // return description IDs
+  return r, nil
+}
+
+// Add feed item and return item ID.
+func (me DbStore) addCve(ctx context.Context, tx Tx, itemId int64, cve nvd_feed.Cve) error {
+  // cve ID
+  var cveId int64
+
+  // check data type, data format, and data version
+  if err := me.checkData(cve.DataType, cve.DataFormat, cve.DataVersion); err != nil {
+    return err
+  }
+
+  // add cve
+  rs, err := tx.Exec(ctx, "feed/insert-cve", itemId, cve.Metadata.Id, cve.Metadata.Assigner)
+  if  err != nil {
+    return err
+  }
+
+  // get cve ID
+  cveId, err = rs.LastInsertId()
+  if err != nil {
+    return err
+  }
+
+  // add problem types
+  for _, pt := range(cve.ProblemTypes.ProblemTypes) {
+    // problem type ID
+    var ptId int64
+
+    // add problem type
+    rs, err := tx.Exec(ctx, "feed/insert-cve-problem", cveId)
+    if err != nil {
+      return err
+    }
+
+    // get problem type ID
+    ptId, err = rs.LastInsertId()
+    if err != nil {
+      return err
+    }
+
+    // add problem type descriptions
+    descIds, err := me.addDescriptions(ctx, tx, pt.Descriptions)
+    if err != nil {
+      return err
+    }
+
+    // add problem type descriptions
+    for _, descId := range(descIds) {
+      _, err = tx.Exec(ctx, "feed/insert-cve-problem-desc", ptId, descId)
+      if err != nil {
+        return err
+      }
+    }
+  }
+
+  // add references
+  for _, ref := range(cve.References.References) {
+    // add reference
+    rs, err := tx.Exec(ctx, "feed/insert-cve-ref", cveId, ref.Url, ref.Name, ref.RefSource)
+    if err != nil {
+      return err
+    }
+
+    // get ref ID
+    refId, err := rs.LastInsertId()
+    if err != nil {
+      return err
+    }
+
+    // add cve ref tags
+    for _, tag := range(ref.Tags) {
+      _, err = tx.Exec(ctx, "feed/insert-cve-ref-tag", refId, tag)
+      if err != nil {
+        return err
+      }
+    }
+  }
+
+  // add cve descriptions
+  descIds, err := me.addDescriptions(ctx, tx, cve.Description.Descriptions)
+  if err != nil {
+    return err
+  }
+
+  // add cve descriptions
+  for _, descId := range(descIds) {
+    _, err = tx.Exec(ctx, "feed/insert-cve-desc", cveId, descId)
+    if err != nil {
+      return err
+    }
+  }
+
+  // return success
+  return nil
+}
+
+// Add feed item and return item ID.
+func (me DbStore) addItem(ctx context.Context, tx Tx, feedId int64, item nvd_feed.Item) error {
+  // item ID
+  var itemId int64
+
+  // serialize published timestamp
+  publishedAt, err := item.PublishedDate.MarshalText()
+  if err != nil {
+    return err
+  }
+
+  // serialize lastmodified timestamp
+  modifiedAt, err := item.LastModifiedDate.MarshalText()
+  if err != nil {
+    return err
+  }
+
+  // add feed item
+  rs, err := tx.Exec(ctx, "feed/insert-item", feedId, publishedAt, modifiedAt)
+  if  err != nil {
+    return err
+  }
+
+  // get item ID
+  itemId, err = rs.LastInsertId()
+  if err != nil {
+    return err
+  }
+
+  // add item cve
+  if err := me.addCve(ctx, tx, itemId, item.Cve); err != nil {
+    return err
+  }
+
+  // TODO: add impact, etc
+
+  // return success
+  return nil
+}
+
+// Add base CVE feed and return feed ID.
+func (me DbStore) addFeed(ctx context.Context, tx Tx, feed nvd_feed.Feed) (int64, error) {
+  // return feed ID
+  var feedId int64
+
+  // check feed data type, data format, and data version
+  if err := me.checkData(feed.DataType, feed.DataFormat, feed.DataVersion); err != nil {
+    return feedId, err
+  }
+
+  // serialize feed timestamp
+  ts, err := feed.Timestamp.MarshalText()
+  if err != nil {
+    return feedId, err
+  }
+
+  // add feed
+  rs, err := tx.Exec(ctx, "feed/insert", string(ts))
+  if err != nil {
+    return feedId, err
+  }
+
+  // get feed ID
+  feedId, err = rs.LastInsertId()
+  if err != nil {
+    return feedId, err
+  }
+
+  // add feed items
+  for _, item := range(feed.Items) {
+    if err := me.addItem(ctx, tx, feedId, item); err != nil {
+      return feedId, err
+    }
+  }
+
+  // return success
+  return feedId, nil
+}
+
+// query IDs used by AddCpeMatches()
+var addCveFeedQueryIds = []string {
+  "feed/insert",
+  "feed/insert-cve",
+}
+
+// Add CVE feed.
+func (me DbStore) AddCveFeed(ctx context.Context, feed nvd_feed.Feed) (int64, error) {
+  // feed ID
+  var feedId int64
+
+  // lazy-init db
+  if err := me.Init(ctx); err != nil {
+    return feedId, err
+  }
+
+  // begin transaction
+  err := me.Tx(ctx, addCveFeedQueryIds, func(tx Tx) error {
+    // add feed, get feed ID
+    if id, err := me.addFeed(ctx, tx, feed); err != nil {
+      return err
+    } else {
+      feedId = id
+      return nil
+    }
+  })
+
+  // return result
+  return feedId, err
+}