move constants to base component

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

Hide file tree

Showing 6 changed files with 49 additions and 53 deletions.
diff --git a/cmd/draconctl/main.go b/cmd/draconctl/main.go
@@ -9,7 +9,7 @@ import (
 	"github.com/ocurity/dracon/cmd/draconctl/components"
 	"github.com/ocurity/dracon/cmd/draconctl/migrations"
 	"github.com/ocurity/dracon/cmd/draconctl/pipelines"
-	draconLogger "github.com/ocurity/dracon/pkg/log"
+	draconlogger "github.com/ocurity/dracon/pkg/log"
 )
 
 var rootCmd = &cobra.Command{
@@ -25,7 +25,7 @@ func main() {
 	pipelines.RegisterPipelinesSubcommands(rootCmd)
 	migrations.RegisterMigrationsSubcommands(rootCmd)
 	components.RegisterComponentsSubcommands(rootCmd)
-	draconLogger.SetDefault(slog.LevelInfo, "", false)
+	draconlogger.SetDefault(slog.LevelInfo, "", false)
 	err := rootCmd.Execute()
 	if err != nil {
 		os.Exit(1)

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
@@ -8,44 +8,35 @@ import (
 	"log/slog"
 	"os"
 
-	apiv1 "github.com/ocurity/dracon/api/proto/v1"
-
+	draconapiv1 "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.
 	Raw bool
-	// Debug flag initializes the logger with a debug level
-	Debug bool
+	// debug flag initializes the logger with a debug level
+	debug bool
 )
 
 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")
-
+	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()
-	if Debug {
-		draconLogger.SetDefault(slog.LevelDebug, os.Getenv(EnvDraconScanID), true)
-	} else {
-		draconLogger.SetDefault(slog.LevelInfo, os.Getenv(EnvDraconScanID), true)
+
+	logLevel := slog.LevelInfo
+	if debug {
+		logLevel = slog.LevelDebug
 	}
+	draconLogger.SetDefault(logLevel, os.Getenv(components.EnvDraconScanID), true)
 
 	if len(inResults) < 1 {
 		return fmt.Errorf("in is undefined")
@@ -54,11 +45,11 @@ func ParseFlags() error {
 }
 
 // LoadToolResponse loads raw results from producers.
-func LoadToolResponse() ([]*apiv1.LaunchToolResponse, error) {
+func LoadToolResponse() ([]*draconapiv1.LaunchToolResponse, error) {
 	return putil.LoadToolResponse(inResults)
 }
 
 // LoadEnrichedToolResponse loads enriched results from the enricher.
-func LoadEnrichedToolResponse() ([]*apiv1.EnrichedLaunchToolResponse, error) {
+func LoadEnrichedToolResponse() ([]*draconapiv1.EnrichedLaunchToolResponse, error) {
 	return putil.LoadEnrichedToolResponse(inResults)
 }
diff --git a/components/consumers/consumer_test.go b/components/consumers/consumer_test.go
@@ -6,11 +6,11 @@ import (
 	"testing"
 	"time"
 
-	v1 "github.com/ocurity/dracon/api/proto/v1"
-
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	draconapiv1 "github.com/ocurity/dracon/api/proto/v1"
+	"github.com/ocurity/dracon/components"
 	"github.com/ocurity/dracon/pkg/putil"
 )
 
@@ -23,7 +23,7 @@ func TestLoadToolResponse(t *testing.T) {
 
 	defer require.NoError(t, os.Remove(tmpFile.Name()))
 
-	issues := []*v1.Issue{
+	issues := []*draconapiv1.Issue{
 		{
 			Target:      "/dracon/source/foobar",
 			Title:       "/dracon/source/barfoo",
@@ -39,9 +39,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
@@ -10,14 +10,14 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"log"
 	"log/slog"
 	"os"
 	"path/filepath"
 	"strings"
 	"time"
 
-	v1 "github.com/ocurity/dracon/api/proto/v1"
+	draconapiv1 "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 +36,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,11 +46,12 @@ func ParseFlags() error {
 	flag.BoolVar(&Append, "append", false, "Append to output file instead of overwriting it")
 
 	flag.Parse()
+	logLevel := slog.LevelInfo
 	if Debug {
-		draconLogger.SetDefault(slog.LevelDebug, os.Getenv(EnvDraconScanID), true)
-	} else {
-		draconLogger.SetDefault(slog.LevelInfo, os.Getenv(EnvDraconScanID), true)
+		logLevel = slog.LevelDebug
 	}
+	draconLogger.SetDefault(logLevel, os.Getenv(components.EnvDraconScanID), true)
+
 	if InResults == "" {
 		return fmt.Errorf("in is undefined")
 	}
@@ -123,28 +117,28 @@ func ParseMultiJSONMessages(in []byte) ([]interface{}, error) {
 // WriteDraconOut provides a generic method to write the resulting protobuf to the output file.
 func WriteDraconOut(
 	toolName string,
-	issues []*v1.Issue,
+	issues []*draconapiv1.Issue,
 ) error {
 	source := getSource()
-	cleanIssues := []*v1.Issue{}
+	cleanIssues := []*draconapiv1.Issue{}
 	for _, iss := range issues {
 		iss.Description = strings.ReplaceAll(iss.Description, sourceDir, ".")
 		iss.Title = strings.ReplaceAll(iss.Title, sourceDir, ".")
 		iss.Target = strings.ReplaceAll(iss.Target, sourceDir, ".")
 		iss.Source = source
 		cleanIssues = append(cleanIssues, iss)
-		log.Printf("found issue: %+v\n", iss)
+		slog.Info(fmt.Sprintf("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 {
-		log.Printf("scan with uuid %s does not have any tags, err:%s", scanUUUID, err)
+		slog.Info(fmt.Sprintf("scan with uuid %s does not have any tags, err:%s", scanUUUID, err))
 	}
 
 	stat, err := os.Stat(OutFile)
@@ -163,7 +157,7 @@ func getSource() string {
 
 	dat, err := os.ReadFile(sourceMetaPath)
 	if err != nil {
-		log.Println(err)
+		slog.Error(err.Error())
 	}
 	return strings.TrimSpace(string(dat))
 }
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