move constants to base component

smithy-security · Jun 4, 2024 · 468b655 · 468b655
1 parent 39bae1e
commit 468b655
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 34 deletions.
diff --git a/components/component.go b/components/component.go
@@ -0,0 +1,10 @@
+package components
+
+const (
+	// EnvDraconStartTime Start Time of Dracon Scan in RFC3339.
+	EnvDraconStartTime = "DRACON_SCAN_TIME"
+	// EnvDraconScanID the ID of the dracon scan.
+	EnvDraconScanID = "DRACON_SCAN_ID"
+	// EnvDraconScanTags the tags of the dracon scan.
+	EnvDraconScanTags = "DRACON_SCAN_TAGS"
+)
diff --git a/components/consumers/consumer.go b/components/consumers/consumer.go
@@ -9,20 +9,12 @@ import (
 	"os"
 
 	apiv1 "github.com/ocurity/dracon/api/proto/v1"
+	"github.com/ocurity/dracon/components"
 
 	draconLogger "github.com/ocurity/dracon/pkg/log"
 	"github.com/ocurity/dracon/pkg/putil"
 )
 
-const (
-	// EnvDraconStartTime Start Time of Dracon Scan in RFC3339.
-	EnvDraconStartTime = "DRACON_SCAN_TIME"
-	// EnvDraconScanID the ID of the dracon scan.
-	EnvDraconScanID = "DRACON_SCAN_ID"
-	// EnvDraconScanTags the tags of the dracon scan.
-	EnvDraconScanTags = "DRACON_SCAN_TAGS"
-)
-
 var (
 	inResults string
 	// Raw represents if the non-enriched results should be used.
@@ -35,18 +27,15 @@ func init() {
 	flag.StringVar(&inResults, "in", "", "the directory where dracon producer/enricher outputs are")
 	flag.BoolVar(&Raw, "raw", false, "if the non-enriched results should be used")
 	flag.BoolVar(&Debug, "debug", false, "turn on debug logging")
-
 }
 
 // ParseFlags will parse the input flags for the consumer and perform simple validation.
 func ParseFlags() error {
 	flag.Parse()
+	draconLogger.SetDefault(slog.LevelInfo, os.Getenv(components.EnvDraconScanID), true)
 	if Debug {
-		draconLogger.SetDefault(slog.LevelDebug, os.Getenv(EnvDraconScanID), true)
-	} else {
-		draconLogger.SetDefault(slog.LevelInfo, os.Getenv(EnvDraconScanID), true)
+		draconLogger.SetDefault(slog.LevelDebug, os.Getenv(components.EnvDraconScanID), true)
 	}
-
 	if len(inResults) < 1 {
 		return fmt.Errorf("in is undefined")
 	}

diff --git a/components/consumers/consumer_test.go b/components/consumers/consumer_test.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	v1 "github.com/ocurity/dracon/api/proto/v1"
+	"github.com/ocurity/dracon/components"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -39,9 +40,9 @@ func TestLoadToolResponse(t *testing.T) {
 	scanTags, err := json.Marshal(tags)
 	assert.NoError(t, err)
 
-	require.NoError(t, os.Setenv(EnvDraconStartTime, timestamp))
-	require.NoError(t, os.Setenv(EnvDraconScanID, scanID))
-	require.NoError(t, os.Setenv(EnvDraconScanTags, string(scanTags)))
+	require.NoError(t, os.Setenv(components.EnvDraconStartTime, timestamp))
+	require.NoError(t, os.Setenv(components.EnvDraconScanID, scanID))
+	require.NoError(t, os.Setenv(components.EnvDraconScanTags, string(scanTags)))
 
 	resultTempDir := tmpFile.Name()
 	resultFile := "test-tool"

diff --git a/components/producers/producer.go b/components/producers/producer.go
@@ -18,6 +18,7 @@ import (
 	"time"
 
 	v1 "github.com/ocurity/dracon/api/proto/v1"
+	"github.com/ocurity/dracon/components"
 
 	draconLogger "github.com/ocurity/dracon/pkg/log"
 	"github.com/ocurity/dracon/pkg/putil"
@@ -36,13 +37,6 @@ var (
 
 const (
 	sourceDir = "/workspace/output"
-
-	// EnvDraconStartTime Start Time of Dracon Scan in RFC3339.
-	EnvDraconStartTime = "DRACON_SCAN_TIME"
-	// EnvDraconScanID the ID of the dracon scan.
-	EnvDraconScanID = "DRACON_SCAN_ID"
-	// EnvDraconScanTags the tags of the dracon scan.
-	EnvDraconScanTags = "DRACON_SCAN_TAGS"
 )
 
 // ParseFlags will parse the input flags for the producer and perform simple validation.
@@ -53,10 +47,9 @@ func ParseFlags() error {
 	flag.BoolVar(&Append, "append", false, "Append to output file instead of overwriting it")
 
 	flag.Parse()
+	draconLogger.SetDefault(slog.LevelInfo, os.Getenv(components.EnvDraconScanID), true)
 	if Debug {
-		draconLogger.SetDefault(slog.LevelDebug, os.Getenv(EnvDraconScanID), true)
-	} else {
-		draconLogger.SetDefault(slog.LevelInfo, os.Getenv(EnvDraconScanID), true)
+		draconLogger.SetDefault(slog.LevelDebug, os.Getenv(components.EnvDraconScanID), true)
 	}
 	if InResults == "" {
 		return fmt.Errorf("in is undefined")
@@ -135,12 +128,12 @@ func WriteDraconOut(
 		cleanIssues = append(cleanIssues, iss)
 		log.Printf("found issue: %+v\n", iss)
 	}
-	scanStartTime := strings.TrimSpace(os.Getenv(EnvDraconStartTime))
+	scanStartTime := strings.TrimSpace(os.Getenv(components.EnvDraconStartTime))
 	if scanStartTime == "" {
 		scanStartTime = time.Now().UTC().Format(time.RFC3339)
 	}
-	scanUUUID := strings.TrimSpace(os.Getenv(EnvDraconScanID))
-	scanTagsStr := strings.TrimSpace(os.Getenv(EnvDraconScanTags))
+	scanUUUID := strings.TrimSpace(os.Getenv(components.EnvDraconScanID))
+	scanTagsStr := strings.TrimSpace(os.Getenv(components.EnvDraconScanTags))
 	scanTags := map[string]string{}
 	err := json.Unmarshal([]byte(scanTagsStr), &scanTags)
 	if err != nil {

diff --git a/components/producers/producer_test.go b/components/producers/producer_test.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	v1 "github.com/ocurity/dracon/api/proto/v1"
+	"github.com/ocurity/dracon/components"
 
 	"github.com/mitchellh/mapstructure"
 	"github.com/stretchr/testify/assert"
@@ -25,8 +26,8 @@ func TestWriteDraconOut(t *testing.T) {
 
 	baseTime := time.Now().UTC()
 	timestamp := baseTime.Format(time.RFC3339)
-	require.NoError(t, os.Setenv(EnvDraconStartTime, timestamp))
-	require.NoError(t, os.Setenv(EnvDraconScanID, "ab3d3290-cd9f-482c-97dc-ec48bdfcc4de"))
+	require.NoError(t, os.Setenv(components.EnvDraconStartTime, timestamp))
+	require.NoError(t, os.Setenv(components.EnvDraconScanID, "ab3d3290-cd9f-482c-97dc-ec48bdfcc4de"))
 
 	OutFile = tmpFile.Name()
 	Append = false
@@ -66,8 +67,8 @@ func TestWriteDraconOutAppend(t *testing.T) {
 
 	baseTime := time.Now().UTC()
 	timestamp := baseTime.Format(time.RFC3339)
-	require.NoError(t, os.Setenv(EnvDraconStartTime, timestamp))
-	require.NoError(t, os.Setenv(EnvDraconScanID, "ab3d3290-cd9f-482c-97dc-ec48bdfcc4de"))
+	require.NoError(t, os.Setenv(components.EnvDraconStartTime, timestamp))
+	require.NoError(t, os.Setenv(components.EnvDraconScanID, "ab3d3290-cd9f-482c-97dc-ec48bdfcc4de"))
 
 	OutFile = tmpFile.Name()
 	Append = true