Skip to content

Commit

Permalink
implement #394 - add a custom key-value pair enricher
Browse files Browse the repository at this point in the history
  • Loading branch information
northdpole committed Oct 3, 2024
1 parent 7e7ef12 commit eee23f0
Show file tree
Hide file tree
Showing 7 changed files with 368 additions and 0 deletions.
5 changes: 5 additions & 0 deletions components/enrichers/custom-annotation/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
# Custom Annotation Enricher

This enricher adds a custom annotation to every finding
Useful for development purposes and for tagging results
Can be used multiple times to add multiple annotations
79 changes: 79 additions & 0 deletions components/enrichers/custom-annotation/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
// Package main of the codeowners enricher
// handles enrichment of individual issues with
// the groups/usernames listed in the github repository
// CODEOWNERS files.
// Owners are matched against the "target" field of the issue
package main

import (
"flag"
"log"
"log/slog"
"strings"

apiv1 "github.com/ocurity/dracon/api/proto/v1"
"github.com/ocurity/dracon/components/enrichers"
)

const defaultAnnotation = ""

var (
annotation string
annotationValue string
)

func enrichIssue(i *apiv1.Issue) (*apiv1.EnrichedIssue, error) {
enrichedIssue := apiv1.EnrichedIssue{}
annotations := map[string]string{
annotation: annotationValue,
}

enrichedIssue = apiv1.EnrichedIssue{
RawIssue: i,
Annotations: annotations,
}
return &enrichedIssue, nil
}

func run() error {
res, err := enrichers.LoadData()
if err != nil {
return err
}
if annotation == "" {
annotation = defaultAnnotation
}
for _, r := range res {
slog.Info("processing results for ", slog.Any("scan", r.ScanInfo))
enrichedIssues := []*apiv1.EnrichedIssue{}
for _, i := range r.GetIssues() {
eI, err := enrichIssue(i)
if err != nil {
slog.Error(err.Error())
continue
}
enrichedIssues = append(enrichedIssues, eI)
}

err := enrichers.WriteData(&apiv1.EnrichedLaunchToolResponse{
OriginalResults: r,
Issues: enrichedIssues,
}, strings.ReplaceAll(annotation, " ", "-"))
if err != nil {
return err
}
}
return nil
}

func main() {
flag.StringVar(&annotation, "annotation", enrichers.LookupEnvOrString("ANNOTATION", defaultAnnotation), "what is the annotation this enricher will add to the issues")
flag.StringVar(&annotationValue, "annotation-value", enrichers.LookupEnvOrString("ANNOTATION_VALUE", defaultAnnotation), "what is the annotation value this enricher will add to the issues")

if err := enrichers.ParseFlags(); err != nil {
log.Fatal(err)
}
if err := run(); err != nil {
log.Fatal(err)
}
}
139 changes: 139 additions & 0 deletions components/enrichers/custom-annotation/main_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,139 @@
package main

import (
"fmt"
"os"
"strings"
"testing"

"github.com/google/go-cmp/cmp"
"github.com/stretchr/testify/require"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/testing/protocmp"

draconv1 "github.com/ocurity/dracon/api/proto/v1"
"github.com/ocurity/dracon/components/enrichers"
)

func TestHandlesZeroFindings(t *testing.T) {
indir, outdir := enrichers.SetupIODirs(t)
mockLaunchToolResponses := enrichers.GetEmptyLaunchToolResponse(t)
for i, r := range mockLaunchToolResponses {
// Write sample enriched responses to indir
encodedProto, err := proto.Marshal(r)
require.NoError(t, err)
rwPermission600 := os.FileMode(0o600)
require.NoError(t, os.WriteFile(fmt.Sprintf("%s/input_%d_%s.tagged.pb", indir, i, r.ToolName), encodedProto, rwPermission600))
}

// Run the enricher
enrichers.SetReadPathForTests(indir)
enrichers.SetWritePathForTests(outdir)
require.NoError(t, run())

// Check there is something in our output directory
files, err := os.ReadDir(outdir)
require.NoError(t, err)
require.NotEmpty(t, files)
require.Len(t, files, 4)

// Check that both of them are EnrichedLaunchToolResponse
// and their Issue property is an empty list
for _, f := range files {
if strings.HasSuffix(f.Name(), ".raw.pb") {
continue
}

encodedProto, err := os.ReadFile(fmt.Sprintf("%s/%s", outdir, f.Name()))
require.NoError(t, err)
output := &draconv1.EnrichedLaunchToolResponse{}
require.NoError(t, proto.Unmarshal(encodedProto, output))
require.Empty(t, output.Issues)
}
}

func TestHandlesFindings(t *testing.T) {
indir, outdir := enrichers.SetupIODirs(t)
annotation = "This is my custom annotation"
annotationValue = "This is my custom annotation value"

mockLaunchToolResponses := enrichers.GetLaunchToolResponse(t)
for i, r := range mockLaunchToolResponses {
// Write sample enriched responses to indir
encodedProto, err := proto.Marshal(r)
require.NoError(t, err)
rwPermission600 := os.FileMode(0o600)
require.NoError(t, os.WriteFile(fmt.Sprintf("%s/input_%d_%s.tagged.pb", indir, i, r.ToolName), encodedProto, rwPermission600))
}

// Run the enricher
enrichers.SetReadPathForTests(indir)
enrichers.SetWritePathForTests(outdir)
require.NoError(t, run())

// Check there is something in our output directory
files, err := os.ReadDir(outdir)
require.NoError(t, err)
require.NotEmpty(t, files)
require.Len(t, files, 4)

// Check that both of them are EnrichedLaunchToolResponse
// and their Issue property is not an empty list
expected := []*draconv1.EnrichedLaunchToolResponse{
{
OriginalResults: mockLaunchToolResponses[0],
Issues: []*draconv1.EnrichedIssue{
{
RawIssue: mockLaunchToolResponses[0].Issues[0],
Annotations: map[string]string{
annotation: annotationValue,
},
},
{
RawIssue: mockLaunchToolResponses[0].Issues[1],
Annotations: map[string]string{
annotation: annotationValue,
},
},
},
},
{
OriginalResults: mockLaunchToolResponses[1],
Issues: []*draconv1.EnrichedIssue{
{
RawIssue: mockLaunchToolResponses[1].Issues[0],
Annotations: map[string]string{
annotation: annotationValue,
},
},
{
RawIssue: mockLaunchToolResponses[1].Issues[1],
Annotations: map[string]string{
annotation: annotationValue,
},
},
},
},
}
var actual draconv1.EnrichedLaunchToolResponse
for _, f := range files {
if strings.HasSuffix(f.Name(), ".raw.pb") {
continue
}
encodedProto, err := os.ReadFile(fmt.Sprintf("%s/%s", outdir, f.Name()))
require.NoError(t, err)

require.NoError(t, proto.Unmarshal(encodedProto, &actual))
if actual.OriginalResults.ToolName == expected[0].OriginalResults.ToolName {
if !proto.Equal(&actual, expected[0]) {
t.Log(cmp.Diff(&actual, expected[0], protocmp.Transform()))
t.Fatal("Actual does not match expected")
}
} else {
if !proto.Equal(&actual, expected[1]) {
t.Log(cmp.Diff(&actual, expected[1], protocmp.Transform()))
t.Fatal("Actual does not match expected")
}
}
}
}
33 changes: 33 additions & 0 deletions components/enrichers/custom-annotation/task.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
---
apiVersion: tekton.dev/v1beta1
kind: Task
metadata:
name: enricher-custom-annotation
labels:
v1.dracon.ocurity.com/component: enricher
spec:
description: Identifies a code owner for each finding.
params:
- name: enricher-custom-annotation-annotation
type: string
default: ""
- name: enricher-custom-annotation-value
type: string
default: ""
workspaces:
- name: output
description: The workspace where we can output results
steps:
- name: run-enricher
imagePullPolicy: IfNotPresent
image: '{{ default "ghcr.io/ocurity/dracon" .Values.image.registry }}/components/enrichers/custom-annotation:{{ .Chart.AppVersion }}'
command: ["/app/components/enrichers/custom-annotation/custom-annotation"]
env:
- name: READ_PATH
value: $(workspaces.output.path)/.dracon/producers
- name: WRITE_PATH
value: "$(workspaces.output.path)/.dracon/enrichers/custom-annotation"
- name: ANNOTATION
value: "$(params.enricher-custom-annotation-annotation)"
- name: ANNOTATION_VALUE
value: "$(params.enricher-custom-annotation-value)"
76 changes: 76 additions & 0 deletions components/enrichers/test_utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,3 +33,79 @@ func GetEmptyLaunchToolResponse(_ *testing.T) []*draconv1.LaunchToolResponse {
},
}
}

// GetEmptyLaunchToolResponse returns a slice of LaunchToolResponse with no issues
func GetLaunchToolResponse(_ *testing.T) []*draconv1.LaunchToolResponse {
code := `this
is
some
code`
return []*draconv1.LaunchToolResponse{
{
ToolName: "tool1",
Issues: []*draconv1.Issue{
{
Target: "file:/a/b/c/d.php:1-2",
Type: "sometype",
Title: "this is a title",
Severity: draconv1.Severity_SEVERITY_CRITICAL,
Cvss: 1.0,
Confidence: draconv1.Confidence_CONFIDENCE_CRITICAL,
Description: "this is a handy dandy description",
Source: "this is a source",
Cve: "CVE-2020-123",
Uuid: "d9681ae9-223b-4df8-a422-7b29bb917a36",
Cwe: []int32{123},
ContextSegment: &code,
},
{
Target: "file:/a/b/c/d.go:2-3",
Type: "sometype1",
Title: "this is a title1",
Severity: draconv1.Severity_SEVERITY_CRITICAL,
Cvss: 1.0,
Confidence: draconv1.Confidence_CONFIDENCE_CRITICAL,
Description: "this is a handy dandy description1",
Source: "this is a source1",
Cve: "CVE-2020-124",
Uuid: "a9681ae9-223b-4df8-a422-7b29bb917a36",
Cwe: []int32{123},
ContextSegment: &code,
},
},
},
{
ToolName: "tool2",
Issues: []*draconv1.Issue{
{
Target: "file:/a/b/c/d.py:1-2",
Type: "sometype",
Title: "this is a title",
Severity: draconv1.Severity_SEVERITY_CRITICAL,
Cvss: 1.0,
Confidence: draconv1.Confidence_CONFIDENCE_CRITICAL,
Description: "this is a handy dandy description",
Source: "this is a source",
Cve: "CVE-2020-123",
Uuid: "q9681ae9-223b-4df8-a422-7b29bb917a36",
Cwe: []int32{123},
ContextSegment: &code,
},
{
Target: "file:/a/b/c/d.py:2-3",
Type: "sometype1",
Title: "this is a title1",
Severity: draconv1.Severity_SEVERITY_CRITICAL,
Cvss: 1.0,
Confidence: draconv1.Confidence_CONFIDENCE_CRITICAL,
Description: "this is a handy dandy description1",
Source: "this is a source1",
Cve: "CVE-2020-124",
Uuid: "w9681ae9-223b-4df8-a422-7b29bb917a36",
Cwe: []int32{123},
ContextSegment: &code,
},
},
},
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
---
apiVersion: kustomize.config.k8s.io/v1beta1
kind: Kustomization
nameSuffix: -annotation-project
components:
- pkg:helm/dracon-oss-components/base
- pkg:helm/dracon-oss-components/git-clone
- pkg:helm/dracon-oss-components/producer-golang-gosec
- pkg:helm/dracon-oss-components/producer-aggregator
- pkg:helm/dracon-oss-components/enricher-custom-annotation
- pkg:helm/dracon-oss-components/enricher-aggregator
- pkg:helm/dracon-oss-components/consumer-stdout-json
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
---
apiVersion: tekton.dev/v1beta1
kind: PipelineRun
metadata:
generateName: dracon-annotation-project-
spec:
pipelineRef:
name: dracon-annotation-project
params:
- name: git-clone-url
value: https://github.com/sqreen/go-dvwa.git
- name: enricher-custom-annotation-annotation
value: "foo"
- name: enricher-custom-annotation-value
value: "bar"
workspaces:
- name: output
volumeClaimTemplate:
spec:
accessModes:
- ReadWriteOnce
resources:
requests:
storage: 1Gi

0 comments on commit eee23f0

Please sign in to comment.