nvdmirror: add SyncConfig, add additional tests

1 files changed, 61 insertions, 40 deletions

diff --git a/nvdmirror/nvdmirror.go b/nvdmirror/nvdmirror.go
index 7b092b1..da758ad 100644
--- a/nvdmirror/nvdmirror.go
+++ b/nvdmirror/nvdmirror.go
@@ -12,28 +12,46 @@ import (
   "time"
 )
 
-// URLs
-type Urls struct {
-  // CVE 1.1 Base URL.  The full meta and JSON URLs are constructed from
-  // this appended to this base..
-  Cve11Base string
+// default user agent (FIXME: make configurable)
+var defaultUserAgent = "cvez/0.1.0"
+
+// Sync() configuration.
+type SyncConfig struct {
+  // CVE 1.1 Base URL.  The full meta and JSON URLs are constructed by
+  // appending the file name to this base.
+  Cve11BaseUrl string
+
+  // CPE 1.0 base URL.  The full meta and JSON URLs are constructed by
+  // appending the file name to this base.
+  Cpe10MatchBaseUrl string
 
-  // CPE 1.0 match meta URL
-  Cpe10MatchMeta string
+  // CPE 2.3 dictionary URL.
+  Cpe23DictUrl string
 
-  // CPE 1.0 match URL
-  Cpe10MatchJson string
+  // User agent string.  Set to "" for default user agent string.
+  UserAgent string
 
-  // CPE 2.3 dictionary URL
-  Cpe23Dict string
+  // Maximum number of idle connections.
+	MaxIdleConns int
+
+  // Idle connection timeout.
+	IdleConnTimeout time.Duration
+}
+
+// Get user agent string.
+func (me SyncConfig) GetUserAgent() string {
+  if len(me.UserAgent) > 0 {
+    return me.UserAgent
+  } else {
+    return defaultUserAgent
+  }
 }
 
 // NVD URLs
-var DefaultUrls = Urls {
-  Cve11Base: "https://nvd.nist.gov/feeds/json/cve/1.1",
-  Cpe10MatchMeta: "https://nvd.nist.gov/feeds/json/cpematch/1.0/nvdcpematch-1.0.meta",
-  Cpe10MatchJson: "https://nvd.nist.gov/feeds/json/cpematch/1.0/nvdcpematch-1.0.gz",
-  Cpe23Dict: "https://nvd.nist.gov/feeds/xml/cpe/dictionary/official-cpe-dictionary_v2.3.xml.gz",
+var DefaultConfig = SyncConfig {
+  Cve11BaseUrl: "https://nvd.nist.gov/feeds/json/cve/1.1",
+  Cpe10MatchBaseUrl: "https://nvd.nist.gov/feeds/json/cpematch/1.0",
+  Cpe23DictUrl: "https://nvd.nist.gov/feeds/xml/cpe/dictionary/official-cpe-dictionary_v2.3.xml.gz",
 }
 
 var extraFiles = []string {
@@ -45,7 +63,7 @@ var extraFiles = []string {
 type fetchResult struct {
   src string // source URL
   err error // fetch result
-  notModified bool // was the result unmodified?
+  modified bool // was the result modified?
   lastModified string // last modified response header
   etag string // etag response header
 }
@@ -53,10 +71,7 @@ type fetchResult struct {
 // base CVE year
 const baseYear = 2002
 
-// default user agent (FIXME: make configurable)
-var userAgent = "cvez/0.1.0"
-
-func fetch(ch chan fetchResult, client *http.Client, cache Cache, srcUrl string, dstDir string) {
+func fetch(ch chan fetchResult, config SyncConfig, cache Cache, client *http.Client, srcUrl string, dstDir string) {
   // parse source url
   src, err := url.Parse(srcUrl)
   if err != nil {
@@ -84,7 +99,7 @@ func fetch(ch chan fetchResult, client *http.Client, cache Cache, srcUrl string,
   }
 
   // add request headers
-  req.Header.Add("user-agent", userAgent)
+  req.Header.Add("user-agent", config.GetUserAgent())
   if headers, ok := cache.Get(srcUrl); ok {
     for k, v := range(headers) {
       req.Header.Add(k, v)
@@ -104,20 +119,21 @@ func fetch(ch chan fetchResult, client *http.Client, cache Cache, srcUrl string,
     // copy body to result
     if size, err := io.Copy(f, resp.Body); err != nil {
       // copy failed
-      ch <- fetchResult { src: srcUrl, err: err }
+      ch <- fetchResult { src: srcUrl, err: err, modified: false }
     } else if err = os.Rename(f.Name(), path); err != nil {
       // rename failed
-      ch <- fetchResult { src: srcUrl, err: err }
+      ch <- fetchResult { src: srcUrl, err: err, modified: false }
     } else {
       log.Debug().Str("url", srcUrl).Int64("size", size).Send()
       ch <- fetchResult {
         src: srcUrl,
+        modified: true,
         lastModified: resp.Header.Get("last-modified"),
         etag: resp.Header.Get("etag"),
       }
     }
   case 304:
-    ch <- fetchResult { src: srcUrl, notModified: true }
+    ch <- fetchResult { src: srcUrl, modified: false }
   default:
     code := resp.StatusCode
     err := fmt.Errorf("%d: %s", code, http.StatusText(code))
@@ -125,7 +141,7 @@ func fetch(ch chan fetchResult, client *http.Client, cache Cache, srcUrl string,
   }
 }
 
-func Sync(urls Urls, cache Cache, dstDir string) error {
+func Sync(config SyncConfig, cache Cache, dstDir string) error {
   log.Debug().Str("dstDir", dstDir).Msg("Sync")
 
   // create fetch result channel
@@ -133,27 +149,27 @@ func Sync(urls Urls, cache Cache, dstDir string) error {
 
   // create shared transport and client
   tr := &http.Transport{
-    // FIXME: make configurable
-	  MaxIdleConns:     10,
-	  IdleConnTimeout:  30 * time.Second,
+	  MaxIdleConns:     config.MaxIdleConns,
+	  IdleConnTimeout:  config.IdleConnTimeout,
   }
   client := &http.Client{Transport: tr}
 
   // calculate total number of years
-  numYears := time.Now().Year() - 2002 + 1
+  numYears := time.Now().Year() - baseYear + 1
 
-  // fetch metas
+  // fetch cve feed metas
   for year := baseYear; year < baseYear + numYears; year++ {
     // build url
-    url := fmt.Sprintf("%s/nvdcve-1.1-%04d.meta", urls.Cve11Base, year)
+    url := fmt.Sprintf("%s/nvdcve-1.1-%04d.meta", config.Cve11BaseUrl, year)
     log.Debug().Int("year", year).Str("url", url).Send()
-    go fetch(ch, client, cache, url, dstDir)
+    go fetch(ch, config, cache, client, url, dstDir)
   }
 
+  // fetch cve extra file metas
   for _, s := range(extraFiles) {
-    url := fmt.Sprintf("%s/%s.meta", urls.Cve11Base, s)
+    url := fmt.Sprintf("%s/%s.meta", config.Cve11BaseUrl, s)
     log.Debug().Str("file", s).Str("url", url).Send()
-    go fetch(ch, client, cache, url, dstDir)
+    go fetch(ch, config, cache, client, url, dstDir)
   }
 
   // read results
@@ -164,14 +180,15 @@ func Sync(urls Urls, cache Cache, dstDir string) error {
     } else {
       log.Info().Str("url", r.src).Str("etag", r.etag).Msg("ok")
 
-      if !r.notModified {
-        // cache headers
-        err := cache.Set(r.src, map[string]string {
+      if r.modified {
+        // build request headers
+        headers := map[string]string {
           "if-none-match": r.etag,
           "if-modified-since": r.lastModified,
-        })
+        }
 
-        if err != nil {
+        // save headers to cache
+        if err := cache.Set(r.src, headers); err != nil {
           log.Error().Str("url", r.src).Err(r.err).Msg("Set")
           return err
         }
@@ -179,6 +196,10 @@ func Sync(urls Urls, cache Cache, dstDir string) error {
     }
   }
 
+  // TODO: fetch cpe dictionary
+  // TODO: fetch cpematch
+  // "nvdcpematch-1.0.{meta,json}",
+
   // return success
   return nil
 }