add util/util.go

1 files changed, 124 insertions, 0 deletions

diff --git a/util/util.go b/util/util.go
new file mode 100644
index 0000000..1c4f84c
--- /dev/null
+++ b/util/util.go
@@ -0,0 +1,124 @@
+// utility functions
+package util
+
+import (
+  "compress/gzip"
+  "encoding/json"
+  "encoding/xml"
+  "github.com/pablotron/cvez/cisa"
+  "github.com/pablotron/cvez/cpedict"
+  "github.com/pablotron/cvez/cpematch"
+  "github.com/pablotron/cvez/nvdmirror"
+  nvd_feed "github.com/pablotron/cvez/feed"
+  "os"
+  "path/filepath"
+)
+
+// read gzipped file, pass reader to callback.
+func gzOpen(path string, fn func(*gzip.Reader) error) error {
+  f, err := os.Open(path)
+  if err != nil {
+    return err
+  }
+  defer f.Close()
+
+  // create gzip reader
+  r, err := gzip.NewReader(f)
+  if err != nil {
+    return err
+  }
+
+  // call func, return result
+  return fn(r)
+}
+
+// read gzipped json, pass decoder to callback
+func jsonGzOpen(path string, fn func(*json.Decoder) error) error {
+  return gzOpen(path, func(r *gzip.Reader) error {
+    return fn(json.NewDecoder(r))
+  })
+}
+
+// read gzipped xml, pass decoder to callback
+func xmlGzOpen(path string, fn func(*xml.Decoder) error) error {
+  return gzOpen(path, func(r *gzip.Reader) error {
+    return fn(xml.NewDecoder(r))
+  })
+}
+
+type getFeedResult struct {
+  feed nvd_feed.Feed
+  err  error
+}
+
+// Read feed from file.
+func getFeed(path string) getFeedResult {
+  var feed nvd_feed.Feed
+
+  err := jsonGzOpen(path, func(d *json.Decoder) error {
+    return d.Decode(&feed)
+  })
+
+  return getFeedResult { feed, err }
+}
+
+// Load CVE feeds in parallel.
+func GetFeeds(cacheDir string, updates []nvdmirror.Update) ([]nvd_feed.Feed, error) {
+  numFeeds := 0
+
+  // load feeds in parallel
+  ch := make(chan getFeedResult)
+  for _, row := range(updates) {
+    if row.Type == nvdmirror.UpdateCveYear {
+      numFeeds += 1
+      go func(path string) {
+        ch <- getFeed(path)
+      }(filepath.Join(cacheDir, row.Path))
+    }
+  }
+
+  // collect feeds
+  feeds := make([]nvd_feed.Feed, numFeeds)
+  for i := 0; i < numFeeds; i++ {
+    if r := <-ch; r.err != nil {
+      return feeds, r.err
+    } else {
+      feeds[i] = r.feed
+    }
+  }
+
+  return feeds, nil
+}
+
+// Read cpe matches from file.
+func GetCpeMatches(path string) (cpematch.Matches, error) {
+  var matches cpematch.Matches
+
+  return matches, jsonGzOpen(path, func(d *json.Decoder) error {
+    return d.Decode(&matches)
+  })
+}
+
+// Read cpe dictionary from file.
+func GetCpeDict(path string) (cpedict.Dictionary, error) {
+  var dict cpedict.Dictionary
+
+  return dict, xmlGzOpen(path, func(d *xml.Decoder) error {
+    return d.Decode(&dict)
+  })
+}
+
+// Read cisa catalog from file.
+func GetCisaCatalog(path string) (cisa.Catalog, error) {
+  var cat cisa.Catalog
+
+  f, err := os.Open(path)
+  if err != nil {
+    return cat, err
+  }
+  defer f.Close()
+
+  // decode/return catalog
+  d := json.NewDecoder(f)
+  return cat, d.Decode(&cat)
+}