epss: enricher should capture all the available data

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

enricher/epss/epss.go

@@ -147,27 +147,34 @@ func (e *Enricher) FetchEnrichment(ctx context.Context, _ driver.Fingerprint) (i
 	var headers []string
 	enc := json.NewEncoder(out)
 	totalCVEs := 0
+	var modelVersion, date string
 
 	for scanner.Scan() {
 		line := strings.TrimSpace(scanner.Text())
-
-		if strings.HasPrefix(line, "#") || line == "" {
-			continue // Skip comment or empty lines
+		if line == "" {
+			continue
+		}
+		// assume metadata is always available at first comment of the file
+		if strings.HasPrefix(line, "#") && date == "" && modelVersion == "" {
+			modelVersion, date = parseMetadata(line)
+			zlog.Info(ctx).
+				Str("modelVersion", modelVersion).
+				Str("scoreDate", date).
+				Msg("parsed metadata")
+			continue
 		}
-
 		if headers == nil {
-			// Store headers on first data line
 			headers = strings.Split(line, ",")
 			continue
 		}
 
 		record := strings.Split(line, ",")
 		if len(record) != len(headers) {
 			zlog.Warn(ctx).Str("line", line).Msg("skipping line with mismatched fields")
-			continue // Skip lines with mismatched number of fields
+			continue
 		}
 
-		r, err := newItemFeed(record, headers)
+		r, err := newItemFeed(record, headers, modelVersion, date)
 		if err != nil {
 			return nil, "", err
 		}
@@ -316,7 +323,7 @@ func (e *Enricher) Enrich(ctx context.Context, g driver.EnrichmentGetter, r *cla
 	return Type, []json.RawMessage{b}, nil
 }
 
-func newItemFeed(record []string, headers []string) (driver.EnrichmentRecord, error) {
+func newItemFeed(record []string, headers []string, modelVersion string, scoreDate string) (driver.EnrichmentRecord, error) {
 	item := make(map[string]interface{}) // Use interface{} to allow mixed types
 	for i, value := range record {
 		// epss details are numeric values
@@ -327,6 +334,13 @@ func newItemFeed(record []string, headers []string) (driver.EnrichmentRecord, er
 		}
 	}
 
+	if modelVersion != "" {
+		item["modelVersion"] = modelVersion
+	}
+	if scoreDate != "" {
+		item["date"] = scoreDate
+	}
+
 	enrichment, err := json.Marshal(item)
 	if err != nil {
 		return driver.EnrichmentRecord{}, fmt.Errorf("failed to encode enrichment: %w", err)
@@ -339,3 +353,28 @@ func newItemFeed(record []string, headers []string) (driver.EnrichmentRecord, er
 
 	return r, nil
 }
+
+func parseMetadata(line string) (modelVersion string, scoreDate string) {
+	// Set default values
+	modelVersion = "N/A"
+	scoreDate = "0001-01-01"
+
+	trimmedLine := strings.TrimPrefix(line, "#")
+	parts := strings.Split(trimmedLine, ",")
+	for _, part := range parts {
+		keyValue := strings.SplitN(part, ":", 2)
+		if len(keyValue) == 2 {
+			key := strings.TrimSpace(keyValue[0])
+			value := strings.TrimSpace(keyValue[1])
+
+			switch key {
+			case "score_date":
+				scoreDate = value
+			case "model_version":
+				modelVersion = value
+			}
+		}
+	}
+
+	return modelVersion, scoreDate
+}

enricher/epss/epss_test.go

@@ -327,23 +327,29 @@ func TestEnrich(t *testing.T) {
 	want := map[string][]map[string]interface{}{
 		"1": {
 			{
-				"cve":        "CVE-2022-34667",
-				"epss":       float64(0.00073),
-				"percentile": float64(0.32799),
+				"cve":          "CVE-2022-34667",
+				"epss":         float64(0.00073),
+				"percentile":   float64(0.32799),
+				"modelVersion": "v2023.03.01",
+				"date":         "2024-10-25T00:00:00+0000",
 			},
 		},
 		"6004": {
 			{
-				"cve":        "CVE-2024-9972",
-				"epss":       float64(0.00091),
-				"percentile": float64(0.39923),
+				"cve":          "CVE-2024-9972",
+				"epss":         float64(0.00091),
+				"percentile":   float64(0.39923),
+				"modelVersion": "v2023.03.01",
+				"date":         "2024-10-25T00:00:00+0000",
 			},
 		},
 		"6005": {
 			{
-				"cve":        "CVE-2024-9986",
-				"epss":       float64(0.00165),
-				"percentile": float64(0.53867),
+				"cve":          "CVE-2024-9986",
+				"epss":         float64(0.00165),
+				"percentile":   float64(0.53867),
+				"modelVersion": "v2023.03.01",
+				"date":         "2024-10-25T00:00:00+0000",
 			},
 		},
 	}
@@ -390,7 +396,7 @@ func parseCSV(filePath string) ([]driver.EnrichmentRecord, error) {
 			continue
 		}
 
-		r, err := newItemFeed(record, headers)
+		r, err := newItemFeed(record, headers, "v2023.03.01", "2024-10-25T00:00:00+0000")
 		if err != nil {
 			return nil, err
 		}