epss: add enrich function

Signed-off-by: daynewlee <[email protected]>

enricher/epss/epss.go

@@ -16,6 +16,8 @@ import (
 	"net/http"
 	"net/url"
 	"path"
+	"regexp"
+	"sort"
 	"strings"
 	"time"
 )
@@ -27,7 +29,17 @@ var (
 	defaultFeed *url.URL
 )
 
+// This is a slightly more relaxed version of the validation pattern in the NVD
+// JSON schema: https://csrc.nist.gov/schema/nvd/feed/1.1/CVE_JSON_4.0_min_1.1.schema
+//
+// It allows for "CVE" to be case insensitive and for dashes and underscores
+// between the different segments.
+var cveRegexp = regexp.MustCompile(`(?i:cve)[-_][0-9]{4}[-_][0-9]{4,}`)
+
 const (
+	// Type is the type of data returned from the Enricher's Enrich method.
+	Type = `message/vnd.clair.map.vulnerability; enricher=clair.epss schema=https://csrc.nist.gov/schema/nvd/feed/1.1/cvss-v3.x.json`
+
 	// DefaultFeeds is the default place to look for EPSS feeds.
 	// epss_scores-YYYY-MM-DD.csv.gz needs to be specified to get all data
 	DefaultFeeds = `https://epss.cyentia.com/`
@@ -136,11 +148,7 @@ func (e *Enricher) FetchEnrichment(ctx context.Context, fingerprint driver.Finge
 		for i, value := range record {
 			item[headers[i]] = value
 		}
-		// filter fields
-		e := make(map[string]string)
-		e["epss"] = item["epss"]
-		e["percentile"] = item["percentile"]
-		enrichment, err := json.Marshal(e)
+		enrichment, err := json.Marshal(item)
 
 		r := driver.EnrichmentRecord{
 			Tags:       []string{item["cve"]},
@@ -198,7 +206,62 @@ func (e *Enricher) sourceURL() {
 	e.feedPath = path.Join(DefaultFeeds, filePath)
 }
 
-func (e *Enricher) Enrich(ctx context.Context, getter driver.EnrichmentGetter, report *claircore.VulnerabilityReport) (string, []json.RawMessage, error) {
-	//TODO implement me
-	panic("implement me")
+func (e *Enricher) Enrich(ctx context.Context, g driver.EnrichmentGetter, r *claircore.VulnerabilityReport) (string, []json.RawMessage, error) {
+	ctx = zlog.ContextWithValues(ctx, "component", "enricher/epss/Enricher/Enrich")
+
+	// We return any CVSS blobs for CVEs mentioned in the free-form parts of the
+	// vulnerability.
+	m := make(map[string][]json.RawMessage)
+
+	erCache := make(map[string][]driver.EnrichmentRecord)
+	for id, v := range r.Vulnerabilities {
+		t := make(map[string]struct{})
+		ctx := zlog.ContextWithValues(ctx,
+			"vuln", v.Name)
+		for _, elem := range []string{
+			v.Description,
+			v.Name,
+			v.Links,
+		} {
+			for _, m := range cveRegexp.FindAllString(elem, -1) {
+				t[m] = struct{}{}
+			}
+		}
+		if len(t) == 0 {
+			continue
+		}
+		ts := make([]string, 0, len(t))
+		for m := range t {
+			ts = append(ts, m)
+		}
+		zlog.Debug(ctx).
+			Strs("cve", ts).
+			Msg("found CVEs")
+
+		sort.Strings(ts)
+		cveKey := strings.Join(ts, "_")
+		rec, ok := erCache[cveKey]
+		if !ok {
+			var err error
+			rec, err = g.GetEnrichment(ctx, ts)
+			if err != nil {
+				return "", nil, err
+			}
+			erCache[cveKey] = rec
+		}
+		zlog.Debug(ctx).
+			Int("count", len(rec)).
+			Msg("found records")
+		for _, r := range rec {
+			m[id] = append(m[id], r.Enrichment)
+		}
+	}
+	if len(m) == 0 {
+		return Type, nil, nil
+	}
+	b, err := json.Marshal(m)
+	if err != nil {
+		return Type, nil, err
+	}
+	return Type, []json.RawMessage{b}, nil
 }