internal/feed: add Vector#MarshalJSON and tests

2 files changed, 42 insertions, 1 deletions

diff --git a/internal/feed/vector.go b/internal/feed/vector.go
index a326c85..9e75de7 100644
--- a/internal/feed/vector.go
+++ b/internal/feed/vector.go
@@ -12,7 +12,7 @@ type Vector struct {
   Vector cvss.Vector
 }
 
-// Unmarshal CVSS score from JSON.
+// Unmarshal CVSS vector from JSON.
 func (me *Vector) UnmarshalJSON(b []byte) error {
   // decode string, check for error
   var s string
@@ -32,3 +32,8 @@ func (me *Vector) UnmarshalJSON(b []byte) error {
   // return success
   return nil
 }
+
+// Marshal CVSS vector to JSON.
+func (me Vector) MarshalJSON() ([]byte, error) {
+  return json.Marshal(me.Vector.String())
+}
diff --git a/internal/feed/vector_test.go b/internal/feed/vector_test.go
index e956884..5dee8ba 100644
--- a/internal/feed/vector_test.go
+++ b/internal/feed/vector_test.go
@@ -2,6 +2,7 @@ package feed
 
 import (
   "encoding/json"
+  "nvd/internal/cvss"
   "testing"
 )
 
@@ -61,3 +62,38 @@ func TestVectorUnmarshalJSON(t *testing.T) {
     })
   }
 }
+
+func TestVectorMarshalJSON(t *testing.T) {
+  tests := []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(tests) {
+    t.Run(val, func(t *testing.T) {
+      // get expected string
+      exp := "\"" + val + "\""
+
+      // create inner vector
+      vec, err := cvss.NewVector(val)
+      if err != nil {
+        t.Error(err)
+        return
+      }
+
+      // serialize as json
+      buf, err := json.Marshal(Vector { vec })
+      if err != nil {
+        t.Error(err)
+        return
+      }
+
+      // check result
+      got := string(buf)
+      if got != exp {
+        t.Errorf("got \"%s\", exp \"%s\"", got, exp)
+      }
+    })
+  }
+}