internal/feed: add vector and tests

3 files changed, 103 insertions, 6 deletions

diff --git a/internal/feed/feed.go b/internal/feed/feed.go
index 7bca496..bdf260c 100644
--- a/internal/feed/feed.go
+++ b/internal/feed/feed.go
@@ -125,10 +125,10 @@ type CvssV3 struct {
   Version V3Version `json:"version"`
 
   // CVSS V3 vector string
-  VectorString string `json:"vectorString"`
+  // VectorString string `json:"vectorString"`
 
-  // CVSS vector (TODO)
-  // Vector cvss.Vector `json:"vectorString"`
+  // CVSS vector
+  Vector Vector `json:"vectorString"`
 
   // attack vector
   AttackVector V3AttackVector `json:"attackVector"`
@@ -171,10 +171,10 @@ type CvssV2 struct {
   Version V2Version `json:"version"`
 
   // CVSS vector string
-  VectorString string `json:"vectorString"`
+  // VectorString string `json:"vectorString"`
 
-  // CVSS vector (TODO)
-  // Vector cvss.Vector `json:"vectorString"`
+  // CVSS vector
+  Vector Vector `json:"vectorString"`
 
   // attack vector
   AccessVector V2AccessVector `json:"accessVector"`
diff --git a/internal/feed/vector.go b/internal/feed/vector.go
new file mode 100644
index 0000000..a326c85
--- /dev/null
+++ b/internal/feed/vector.go
@@ -0,0 +1,34 @@
+// NVD JSON feed parser.
+package feed
+
+import (
+  "encoding/json"
+  "nvd/internal/cvss"
+)
+
+// CVSS vector
+type Vector struct {
+  // CVSS vector
+  Vector cvss.Vector
+}
+
+// Unmarshal CVSS score from JSON.
+func (me *Vector) UnmarshalJSON(b []byte) error {
+  // decode string, check for error
+  var s string
+  if err := json.Unmarshal(b, &s); err != nil {
+    return err
+  }
+
+  // parse vector
+  vec, err := cvss.NewVector(s)
+  if err != nil {
+    return err
+  }
+
+  // save result
+  me.Vector = vec
+
+  // return success
+  return nil
+}
diff --git a/internal/feed/vector_test.go b/internal/feed/vector_test.go
new file mode 100644
index 0000000..e956884
--- /dev/null
+++ b/internal/feed/vector_test.go
@@ -0,0 +1,63 @@
+package feed
+
+import (
+  "encoding/json"
+  "testing"
+)
+
+func TestVectorUnmarshalInvalidData(t *testing.T) {
+  test := []byte(`{}`)
+  var val Vector
+
+  if err := json.Unmarshal(test, &val); err == nil {
+    t.Errorf("got \"%s\", exp error", val)
+  }
+}
+
+func TestVectorUnmarshalJSON(t *testing.T) {
+  failTests := []struct {
+    val string
+    exp string
+  } {
+    {
+      val: "\"AV:N/junk/Au:S/C:P/I:P/A:P\"",
+      exp: "invalid CVSS vector: AV:N/junk/Au:S/C:P/I:P/A:P",
+    }, {
+      val: "\"CVSS:3.0/junk/AC:H/PR:H/UI:R/S:U/C:H/I:H/A:H\"",
+      exp: "invalid CVSS vector: CVSS:3.0/junk/AC:H/PR:H/UI:R/S:U/C:H/I:H/A:H",
+    }, {
+      val: "\"CVSS:3.1/AV:A/junk/PR:N/UI:N/S:U/C:H/I:H/A:H\"",
+      exp: "invalid CVSS vector: CVSS:3.1/AV:A/junk/PR:N/UI:N/S:U/C:H/I:H/A:H",
+    },
+  }
+
+  for _, test := range(failTests) {
+    t.Run(test.val, func(t *testing.T) {
+      var got Vector
+
+      if err := json.Unmarshal([]byte(test.val), &got); err == nil {
+        t.Errorf("got \"%v\", exp error", got)
+      } else if err.Error() != test.exp {
+        t.Errorf("got \"%s\", exp \"%s\"", err.Error(), test.exp)
+      }
+    })
+  }
+
+  passTests := []string {
+    "AV:N/AC:M/Au:S/C:P/I:P/A:P",
+    "CVSS:3.0/AV:A/AC:H/PR:H/UI:R/S:U/C:H/I:H/A:H",
+    "CVSS:3.1/AV:A/AC:H/PR:N/UI:N/S:U/C:H/I:H/A:H",
+  }
+
+  for _, val := range(passTests) {
+    t.Run(val, func(t *testing.T) {
+      var got Vector
+
+      if err := json.Unmarshal([]byte("\"" + val + "\""), &got); err != nil {
+        t.Error(err)
+      } else if got.Vector.String() != val {
+        t.Errorf("got \"%s\", exp \"%s\"", got.Vector.String(), val)
+      }
+    })
+  }
+}