Skip to content

Commit

Permalink
move constants to base component
Browse files Browse the repository at this point in the history
  • Loading branch information
northdpole authored and ptzianos committed Jun 5, 2024
1 parent e248a41 commit db05659
Show file tree
Hide file tree
Showing 7 changed files with 52 additions and 78 deletions.
3 changes: 1 addition & 2 deletions cmd/draconctl/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ 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"
)

var rootCmd = &cobra.Command{
Expand All @@ -25,7 +24,7 @@ func main() {
pipelines.RegisterPipelinesSubcommands(rootCmd)
migrations.RegisterMigrationsSubcommands(rootCmd)
components.RegisterComponentsSubcommands(rootCmd)
draconLogger.SetDefault(slog.LevelInfo, "", false)
slog.SetDefault(slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{})))
err := rootCmd.Execute()
if err != nil {
os.Exit(1)
Expand Down
10 changes: 10 additions & 0 deletions components/component.go
Original file line number Diff line number Diff line change
@@ -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"
)
34 changes: 12 additions & 22 deletions components/consumers/consumer.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,57 +8,47 @@ import (
"log/slog"
"os"

apiv1 "github.com/ocurity/dracon/api/proto/v1"

draconLogger "github.com/ocurity/dracon/pkg/log"
draconapiv1 "github.com/ocurity/dracon/api/proto/v1"
"github.com/ocurity/dracon/components"
"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
}

slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: logLevel})).With("scanID", os.Getenv(components.EnvDraconScanID)))
if len(inResults) < 1 {
return fmt.Errorf("in is undefined")
}
return nil
}

// 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)
}
12 changes: 6 additions & 6 deletions components/consumers/consumer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand All @@ -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",
Expand All @@ -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"
Expand Down
43 changes: 18 additions & 25 deletions components/producers/producer.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,15 @@ 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"
)

Expand All @@ -30,34 +29,28 @@ var (
OutFile string
// Append flag will append to the outfile instead of overwriting, useful when there's multiple inresults.
Append bool
// Debug flag initializes the logger with a debug level
Debug bool
// debug flag initializes the logger with a debug level
debug bool
)

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.
func ParseFlags() error {
flag.StringVar(&InResults, "in", "", "")
flag.StringVar(&OutFile, "out", "", "")
flag.BoolVar(&Debug, "debug", false, "turn on debug logging")
flag.BoolVar(&debug, "debug", false, "turn on debug logging")
flag.BoolVar(&Append, "append", false, "Append to output file instead of overwriting it")

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
}
slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: logLevel})).With("scanID", os.Getenv(components.EnvDraconScanID)))

if InResults == "" {
return fmt.Errorf("in is undefined")
}
Expand Down Expand Up @@ -123,28 +116,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.Error(fmt.Sprintf("scan does not have any tags, err:%s", err))
}

stat, err := os.Stat(OutFile)
Expand All @@ -163,7 +156,7 @@ func getSource() string {

dat, err := os.ReadFile(sourceMetaPath)
if err != nil {
log.Println(err)
slog.Error(err.Error())
}
return strings.TrimSpace(string(dat))
}
9 changes: 5 additions & 4 deletions components/producers/producer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
19 changes: 0 additions & 19 deletions pkg/log/dracon-logger.go

This file was deleted.

0 comments on commit db05659

Please sign in to comment.