add tests for the base enricher

components/consumers/consumer_test.go

@@ -30,7 +30,7 @@ func TestLoadToolResponse(t *testing.T) {
 			Description: "/dracon/source/example.yaml",
 		},
 	}
-	timestamp := time.Now().UTC().Format(time.RFC3339)
+	timestamp := time.Now().UTC()
 	scanID := "ab3d3290-cd9f-482c-97dc-ec48bdfcc4de"
 	tags := map[string]string{
 		"assetID":       "someID",
@@ -39,7 +39,7 @@ func TestLoadToolResponse(t *testing.T) {
 	scanTags, err := json.Marshal(tags)
 	assert.NoError(t, err)
 
-	require.NoError(t, os.Setenv(components.EnvDraconStartTime, timestamp))
+	require.NoError(t, os.Setenv(components.EnvDraconStartTime, timestamp.Format(time.RFC3339)))
 	require.NoError(t, os.Setenv(components.EnvDraconScanID, scanID))
 	require.NoError(t, os.Setenv(components.EnvDraconScanTags, string(scanTags)))
 

components/enrichers/codeowners/main.go

@@ -16,6 +16,7 @@ import (
 	owners "github.com/hairyhenderson/go-codeowners"
 
 	apiv1 "github.com/ocurity/dracon/api/proto/v1"
+	v1 "github.com/ocurity/dracon/api/proto/v1"
 	"github.com/ocurity/dracon/components/enrichers"
 )
 
@@ -80,7 +81,10 @@ func run() error {
 			}
 			enrichedIssues = append(enrichedIssues, eI)
 		}
-		return enrichers.WriteData(enrichedIssues, r, "codeowners")
+		return enrichers.WriteData(&v1.EnrichedLaunchToolResponse{
+			OriginalResults: r,
+			Issues:          enrichedIssues,
+		}, "codeowners")
 	}
 	return nil
 }

components/enrichers/deduplication/main.go

@@ -53,7 +53,10 @@ func run() error {
 			enrichedIssues = append(enrichedIssues, eI)
 			log.Printf("enriched issue '%s'", eI.GetRawIssue().GetUuid())
 		}
-		return enrichers.WriteData(enrichedIssues, r, "deduplication")
+		return enrichers.WriteData(&v1.EnrichedLaunchToolResponse{
+			OriginalResults: r,
+			Issues:          enrichedIssues,
+		}, "deduplication")
 	}
 	return nil
 }

components/enrichers/depsdev/main.go

@@ -283,7 +283,10 @@ func run() error {
 			}
 			enrichedIssues = append(enrichedIssues, eI)
 		}
-		return enrichers.WriteData(enrichedIssues, r, "deps-dev")
+		return enrichers.WriteData(&v1.EnrichedLaunchToolResponse{
+			OriginalResults: r,
+			Issues:          enrichedIssues,
+		}, "deps-dev")
 	}
 	return nil
 }

components/enrichers/enricher.go

@@ -4,15 +4,14 @@ package enrichers
 import (
 	"flag"
 	"fmt"
-	"log"
 	"log/slog"
 	"os"
 	"path/filepath"
-	"time"
 
 	"github.com/go-errors/errors"
 
 	draconapiv1 "github.com/ocurity/dracon/api/proto/v1"
+	v1 "github.com/ocurity/dracon/api/proto/v1"
 	"github.com/ocurity/dracon/components"
 	"github.com/ocurity/dracon/pkg/putil"
 )
@@ -56,28 +55,29 @@ func LoadData() ([]*draconapiv1.LaunchToolResponse, error) {
 	return putil.LoadTaggedToolResponse(readPath)
 }
 
-func WriteData(enrichedIssues []*draconapiv1.EnrichedIssue, originalResults *draconapiv1.LaunchToolResponse, enricherName string) error {
-	if len(enrichedIssues) > 0 {
-		if err := putil.WriteEnrichedResults(originalResults, enrichedIssues,
-			filepath.Join(writePath, fmt.Sprintf("%s.%s.enriched.pb", originalResults.GetToolName(), enricherName)),
-		); err != nil {
-			return err
-		}
-	} else {
-		log.Println("no enriched issues were created for", originalResults.GetToolName())
+func WriteData(enrichedLaunchToolResponse *v1.EnrichedLaunchToolResponse, enricherName string) error {
+
+	if enrichedLaunchToolResponse == nil || len(enrichedLaunchToolResponse.Issues) == 0 {
+		return errors.Errorf("no enriched issues were created for %s", enrichedLaunchToolResponse.GetOriginalResults().GetToolName())
+	}
+	if err := putil.WriteEnrichedResults(enrichedLaunchToolResponse.GetOriginalResults(), enrichedLaunchToolResponse.GetIssues(),
+		filepath.Join(writePath, fmt.Sprintf("%s.%s.enriched.pb", enrichedLaunchToolResponse.GetOriginalResults().GetToolName(), enricherName)),
+	); err != nil {
+		return err
+	}
+	if enrichedLaunchToolResponse.OriginalResults == nil || len(enrichedLaunchToolResponse.OriginalResults.GetIssues()) == 0 {
+		return errors.Errorf("original results is empty for %s", enrichedLaunchToolResponse.GetOriginalResults().GetToolName())
 	}
-	if len(originalResults.GetIssues()) > 0 {
-		scanStartTime := originalResults.GetScanInfo().GetScanStartTime().AsTime()
-		if err := putil.WriteResults(
-			originalResults.GetToolName(),
-			originalResults.GetIssues(),
-			filepath.Join(writePath, fmt.Sprintf("%s.raw.pb", originalResults.GetToolName())),
-			originalResults.GetScanInfo().GetScanUuid(),
-			scanStartTime.Format(time.RFC3339),
-			originalResults.GetScanInfo().GetScanTags(),
-		); err != nil {
-			return errors.Errorf("could not write results: %s", err)
-		}
+	scanStartTime := enrichedLaunchToolResponse.GetOriginalResults().GetScanInfo().GetScanStartTime().AsTime()
+	if err := putil.WriteResults(
+		enrichedLaunchToolResponse.GetOriginalResults().GetToolName(),
+		enrichedLaunchToolResponse.GetOriginalResults().GetIssues(),
+		filepath.Join(writePath, fmt.Sprintf("%s.raw.pb", enrichedLaunchToolResponse.GetOriginalResults().GetToolName())),
+		enrichedLaunchToolResponse.GetOriginalResults().GetScanInfo().GetScanUuid(),
+		scanStartTime,
+		enrichedLaunchToolResponse.GetOriginalResults().GetScanInfo().GetScanTags(),
+	); err != nil {
+		return errors.Errorf("could not write results: %s", err)
 	}
 	return nil
 }

components/enrichers/enricher_test.go

@@ -0,0 +1,107 @@
+package enrichers
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/testing/protocmp"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	draconapiv1 "github.com/ocurity/dracon/api/proto/v1"
+	v1 "github.com/ocurity/dracon/api/proto/v1"
+	"github.com/ocurity/dracon/pkg/putil"
+)
+
+func createObjects() *draconapiv1.EnrichedLaunchToolResponse {
+	scanID := "7c78f6c9-b4b0-493c-a912-0bb0a4aaaaa0"
+	times, _ := time.Parse(time.RFC3339, "2023-01-19T18:09:06.370037788Z")
+	timestamp := timestamppb.New(times)
+	si := v1.ScanInfo{
+		ScanUuid:      scanID,
+		ScanStartTime: timestamp,
+	}
+	toolName := "SAT-Tool"
+	response := v1.LaunchToolResponse{
+		ToolName: toolName,
+		ScanInfo: &si,
+	}
+	enrichedResponse := v1.EnrichedLaunchToolResponse{}
+
+	var issues []*v1.Issue
+	var enrichedIssues []*v1.EnrichedIssue
+	for j := 0; j < 10; j++ {
+		id := uuid.New()
+		x := v1.Issue{
+			Target:     fmt.Sprintf("target-%d", j),
+			Type:       fmt.Sprintf("type-%d", j),
+			Title:      fmt.Sprintf("title-%d", j),
+			Severity:   v1.Severity_SEVERITY_INFO,
+			Confidence: v1.Confidence_CONFIDENCE_INFO,
+			Uuid:       id.String(),
+		}
+		y := v1.EnrichedIssue{
+			RawIssue:      &x,
+			FirstSeen:     response.ScanInfo.ScanStartTime,
+			Count:         uint64(j),
+			FalsePositive: false,
+			UpdatedAt:     response.ScanInfo.ScanStartTime,
+			Hash:          fmt.Sprintf("d41d8cd98f00b204e9800998ecf842%d", j),
+		}
+		issues = append(issues, &x)
+		enrichedIssues = append(enrichedIssues, &y)
+	}
+	response.Issues = issues
+	enrichedResponse.OriginalResults = &response
+	enrichedResponse.Issues = enrichedIssues
+	return &enrichedResponse
+}
+
+func TestWriteDataNormalOperation(t *testing.T) {
+	enricherName := "tests-enricher"
+	// prepare
+	workdir, err := os.MkdirTemp("/tmp", "")
+	require.NoError(t, err)
+	require.NoError(t, os.Mkdir(filepath.Join(workdir, "raw"), 0755))
+
+	// test errors first
+	enrichedResponse := createObjects()
+	require.Error(t, WriteData(nil, enricherName))
+
+	enrichedResponse.Issues = []*v1.EnrichedIssue{}
+	require.Error(t, WriteData(enrichedResponse, enricherName))
+
+	enrichedResponse = createObjects() // reset
+
+	enrichedResponse.OriginalResults.Issues = []*v1.Issue{}
+	require.Error(t, WriteData(enrichedResponse, enricherName))
+
+	enrichedResponse.OriginalResults = nil
+	require.Error(t, WriteData(enrichedResponse, enricherName))
+
+	// happy path
+	enrichedResponse = createObjects()
+	SetWritePathForTests(workdir)
+	require.NoError(t, WriteData(enrichedResponse, enricherName))
+
+	require.NoError(t, err)
+	er, err := putil.LoadEnrichedNonAggregatedToolResponse(workdir)
+	require.NoError(t, err)
+	opt := cmp.Comparer(func(x, y timestamppb.Timestamp) bool {
+		return x.Nanos == y.Nanos
+	})
+
+	require.True(t, cmp.Equal([]*v1.EnrichedLaunchToolResponse{enrichedResponse}, er, protocmp.Transform(), opt),
+		cmp.Diff([]*v1.EnrichedLaunchToolResponse{enrichedResponse}, er, protocmp.Transform()))
+
+	r, err := putil.LoadToolResponse(filepath.Join(workdir, fmt.Sprintf("%s.raw.pb", enrichedResponse.GetOriginalResults().GetToolName())))
+	require.NoError(t, err)
+
+	require.True(t, cmp.Equal([]*v1.LaunchToolResponse{enrichedResponse.OriginalResults}, r, protocmp.Transform(), opt),
+		cmp.Diff([]*v1.LaunchToolResponse{enrichedResponse.OriginalResults}, r, protocmp.Transform()))
+}

components/enrichers/policy/main.go

@@ -78,7 +78,11 @@ func run() error {
 			}
 			enrichedIssues = append(enrichedIssues, eI)
 		}
-		return enrichers.WriteData(enrichedIssues, r, "policy")
+
+		return enrichers.WriteData(&v1.EnrichedLaunchToolResponse{
+			OriginalResults: r,
+			Issues:          enrichedIssues,
+		}, "policy")
 	}
 	return nil
 }

components/producers/aggregator/main.go

@@ -6,7 +6,6 @@ import (
 	"log"
 	"os"
 	"path/filepath"
-	"time"
 
 	"github.com/google/uuid"
 
@@ -47,13 +46,12 @@ func run() {
 			}
 			taggedIssues = append(taggedIssues, eI)
 		}
-		scanStartTime := r.GetScanInfo().GetScanStartTime().AsTime()
 		if err := putil.WriteResults(
 			r.GetToolName(),
 			taggedIssues,
 			filepath.Join(writePath, fmt.Sprintf("%s.tagged.pb", r.GetToolName())),
 			r.GetScanInfo().GetScanUuid(),
-			scanStartTime.Format(time.RFC3339),
+			r.GetScanInfo().GetScanStartTime().AsTime(),
 			r.GetScanInfo().GetScanTags(),
 		); err != nil {
 			log.Fatalf("could not write results: %s", err)

components/producers/producer.go

@@ -128,14 +128,14 @@ func WriteDraconOut(
 		cleanIssues = append(cleanIssues, iss)
 		slog.Info(fmt.Sprintf("found issue: %+v\n", iss))
 	}
-	scanStartTime := strings.TrimSpace(os.Getenv(components.EnvDraconStartTime))
-	if scanStartTime == "" {
-		scanStartTime = time.Now().UTC().Format(time.RFC3339)
+	scanStartTime, err := time.Parse(time.RFC3339, strings.TrimSpace(os.Getenv(components.EnvDraconStartTime)))
+	if err != nil {
+		scanStartTime = time.Now().UTC()
 	}
 	scanUUUID := strings.TrimSpace(os.Getenv(components.EnvDraconScanID))
 	scanTagsStr := strings.TrimSpace(os.Getenv(components.EnvDraconScanTags))
 	scanTags := map[string]string{}
-	err := json.Unmarshal([]byte(scanTagsStr), &scanTags)
+	err = json.Unmarshal([]byte(scanTagsStr), &scanTags)
 	if err != nil {
 		slog.Error(fmt.Sprintf("scan does not have any tags, err:%s", err))
 	}

go.mod

@@ -29,7 +29,7 @@ require (
 	golang.org/x/crypto v0.18.0
 	golang.org/x/oauth2 v0.16.0
 	google.golang.org/api v0.156.0
-	google.golang.org/protobuf v1.33.0
+	google.golang.org/protobuf v1.34.1
 	gopkg.in/h2non/gock.v1 v1.1.2
 	helm.sh/helm/v3 v3.14.4
 	k8s.io/api v0.29.0

go.sum

@@ -1805,6 +1805,8 @@ google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQ
 google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
 google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
+google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg=
+google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
 gopkg.in/airbrake/gobrake.v2 v2.0.9/go.mod h1:/h5ZAUhDkGaJfjzjKLSjv6zCL6O0LLBxU4K+aSYdM/U=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

pkg/putil/write.go

@@ -44,17 +44,13 @@ func WriteResults(
 	issues []*v1.Issue,
 	outFile string,
 	scanUUID string,
-	scanStartTime string,
+	scanStartTime time.Time,
 	scanTags map[string]string,
 ) error {
 	if err := os.MkdirAll(filepath.Dir(outFile), os.ModePerm); err != nil {
 		return err
 	}
-	timeVal, err := time.Parse(time.RFC3339, scanStartTime)
-	if err != nil {
-		return err
-	}
-	timestamp := timestamppb.New(timeVal)
+	timestamp := timestamppb.New(scanStartTime)
 
 	scanInfo := v1.ScanInfo{
 		ScanUuid:      scanUUID,