fix(analyze): enable sigs consuming sigs

Signatures consuming other signatures relied on the single binary reprocessing finding events into the pipeline. This did not occur in analyze mode. Introduce that mechanism so that it now works. Co-authored-by: Asaf Eitani <asaf.eitani@aquasec.com>
aquasecurity · Nov 5, 2024 · 1d99241 · 1d99241
1 parent 171dccf
commit 1d99241
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 94 deletions.
diff --git a/cmd/tracee/cmd/analyze.go b/cmd/tracee/cmd/analyze.go
@@ -14,8 +14,8 @@ import (
 
 	"github.com/aquasecurity/tracee/pkg/cmd/flags"
 	"github.com/aquasecurity/tracee/pkg/cmd/initialize/sigs"
-	tracee "github.com/aquasecurity/tracee/pkg/ebpf"
 	"github.com/aquasecurity/tracee/pkg/events"
+	"github.com/aquasecurity/tracee/pkg/events/findings"
 	"github.com/aquasecurity/tracee/pkg/logger"
 	"github.com/aquasecurity/tracee/pkg/signatures/engine"
 	"github.com/aquasecurity/tracee/pkg/signatures/signature"
@@ -77,127 +77,141 @@ tracee analyze --events anti_debugging events.json`,
 		bindViperFlag(cmd, "rego")
 		bindViperFlag(cmd, "signatures-dir")
 	},
-	Run: func(cmd *cobra.Command, args []string) {
-		logFlags := viper.GetStringSlice("log")
+	Run:                   command,
+	DisableFlagsInUseLine: true,
+}
 
-		logCfg, err := flags.PrepareLogger(logFlags, true)
-		if err != nil {
-			logger.Fatalw("Failed to prepare logger", "error", err)
-		}
-		logger.Init(logCfg)
+func command(cmd *cobra.Command, args []string) {
+	logFlags := viper.GetStringSlice("log")
 
-		inputFile, err := os.Open(args[0])
-		if err != nil {
-			logger.Fatalw("Failed to get signatures-dir flag", "err", err)
-		}
+	logCfg, err := flags.PrepareLogger(logFlags, true)
+	if err != nil {
+		logger.Fatalw("Failed to prepare logger", "error", err)
+	}
+	logger.Init(logCfg)
 
-		// Rego command line flags
+	inputFile, err := os.Open(args[0])
+	if err != nil {
+		logger.Fatalw("Failed to get signatures-dir flag", "err", err)
+	}
 
-		rego, err := flags.PrepareRego(viper.GetStringSlice("rego"))
-		if err != nil {
-			logger.Fatalw("Failed to parse rego flags", "err", err)
-		}
+	// Rego command line flags
 
-		// Signature directory command line flags
+	rego, err := flags.PrepareRego(viper.GetStringSlice("rego"))
+	if err != nil {
+		logger.Fatalw("Failed to parse rego flags", "err", err)
+	}
 
-		signatureEvents := viper.GetStringSlice("events")
-		// if no event was passed, load all events
-		if len(signatureEvents) == 0 {
-			signatureEvents = nil
-		}
+	// Signature directory command line flags
 
-		signatures, _, err := signature.Find(
-			rego.RuntimeTarget,
-			rego.PartialEval,
-			viper.GetStringSlice("signatures-dir"),
-			signatureEvents,
-			rego.AIO,
-		)
+	signatureEvents := viper.GetStringSlice("events")
+	// if no event was passed, load all events
+	if len(signatureEvents) == 0 {
+		signatureEvents = nil
+	}
 
-		if err != nil {
-			logger.Fatalw("Failed to find signature event", "err", err)
-		}
+	signatures, _, err := signature.Find(
+		rego.RuntimeTarget,
+		rego.PartialEval,
+		viper.GetStringSlice("signatures-dir"),
+		signatureEvents,
+		rego.AIO,
+	)
 
-		if len(signatures) == 0 {
-			logger.Fatalw("No signature event loaded")
-		}
+	if err != nil {
+		logger.Fatalw("Failed to find signature event", "err", err)
+	}
 
-		logger.Infow(
-			"Signatures loaded",
-			"total", len(signatures),
-			"signatures", getSigsNames(signatures),
-		)
+	if len(signatures) == 0 {
+		logger.Fatalw("No signature event loaded")
+	}
 
-		_ = sigs.CreateEventsFromSignatures(events.StartSignatureID, signatures)
+	logger.Infow(
+		"Signatures loaded",
+		"total", len(signatures),
+		"signatures", getSigsNames(signatures),
+	)
 
-		engineConfig := engine.Config{
-			Signatures:          signatures,
-			SignatureBufferSize: 1000,
-		}
+	_ = sigs.CreateEventsFromSignatures(events.StartSignatureID, signatures)
 
-		ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
-		defer stop()
+	engineConfig := engine.Config{
+		Signatures:          signatures,
+		SignatureBufferSize: 1000,
+	}
 
-		engineOutput := make(chan *detect.Finding)
-		engineInput := make(chan protocol.Event)
+	// two seperate contexts.
+	// 1. signal notifiable context that can terminate both analyze and engine work
+	// 2. signal solely to notify internally inside analyze once file input is over
+	signalCtx, _ := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	fileReadCtx, stop := context.WithCancel(signalCtx)
 
-		source := engine.EventSources{Tracee: engineInput}
-		sigEngine, err := engine.NewEngine(engineConfig, source, engineOutput)
-		if err != nil {
-			logger.Fatalw("Failed to create engine", "err", err)
-		}
+	engineOutput := make(chan *detect.Finding)
+	engineInput := make(chan protocol.Event)
 
-		err = sigEngine.Init()
-		if err != nil {
-			logger.Fatalw("failed to initialize signature engine", "err", err)
-		}
+	source := engine.EventSources{Tracee: engineInput}
+	sigEngine, err := engine.NewEngine(engineConfig, source, engineOutput)
+	if err != nil {
+		logger.Fatalw("Failed to create engine", "err", err)
+	}
 
-		go sigEngine.Start(ctx)
+	err = sigEngine.Init()
+	if err != nil {
+		logger.Fatalw("failed to initialize signature engine", "err", err)
+	}
 
-		// producer
-		go produce(ctx, inputFile, engineInput)
+	go sigEngine.Start(signalCtx)
 
-		// consumer
-		for {
-			select {
-			case finding, ok := <-engineOutput:
-				if !ok {
-					return
-				}
-				process(finding)
-			case <-ctx.Done():
-				goto drain
+	// producer
+	go produce(fileReadCtx, stop, inputFile, engineInput)
+
+	// consumer
+	for {
+		select {
+		case finding, ok := <-engineOutput:
+			if !ok {
+				return
 			}
+			process(finding)
+		case <-fileReadCtx.Done():
+			// ensure the engineInput channel will be closed
+			goto drain
+		case <-signalCtx.Done():
+			// ensure the engineInput channel will be closed
+			goto drain
 		}
-	drain:
-		// drain
-		for {
-			select {
-			case finding, ok := <-engineOutput:
-				if !ok {
-					return
-				}
-				process(finding)
-			default:
+	}
+drain:
+	// drain
+	defer close(engineInput)
+	for {
+		select {
+		case finding, ok := <-engineOutput:
+			if !ok {
 				return
 			}
+
+			process(finding)
+		default:
+			return
 		}
-	},
-	DisableFlagsInUseLine: true,
+	}
 }
 
-func produce(ctx context.Context, inputFile *os.File, engineInput chan protocol.Event) {
-	// ensure the engineInput channel will be closed
-	defer close(engineInput)
-
+func produce(ctx context.Context, cancel context.CancelFunc, inputFile *os.File, engineInput chan<- protocol.Event) {
 	scanner := bufio.NewScanner(inputFile)
 	scanner.Split(bufio.ScanLines)
 	for {
 		select {
 		case <-ctx.Done():
+			// if terminated from above
 			return
 		default:
 			if !scanner.Scan() { // if EOF or error close the done channel and return
+				if err := scanner.Err(); err != nil {
+					logger.Errorw("Error while scanning input file", "error", err)
+				}
+				// terminate analysis here and proceed to draining
+				cancel()
 				return
 			}
 
@@ -212,7 +226,7 @@ func produce(ctx context.Context, inputFile *os.File, engineInput chan protocol.
 }
 
 func process(finding *detect.Finding) {
-	event, err := tracee.FindingToEvent(finding)
+	event, err := findings.FindingToEvent(finding)
 	if err != nil {
 		logger.Fatalw("Failed to convert finding to event", "err", err)
 	}
@@ -232,15 +246,15 @@ func bindViperFlag(cmd *cobra.Command, flag string) {
 	}
 }
 
-func getSigsNames(sigs []detect.Signature) []string {
-	var sigsNames []string
-	for _, sig := range sigs {
+func getSigsNames(signatures []detect.Signature) []string {
+	var sigNames []string
+	for _, sig := range signatures {
 		sigMeta, err := sig.GetMetadata()
 		if err != nil {
 			logger.Warnw("Failed to get signature metadata", "err", err)
 			continue
 		}
-		sigsNames = append(sigsNames, sigMeta.Name)
+		sigNames = append(sigNames, sigMeta.Name)
 	}
-	return sigsNames
+	return sigNames
 }
diff --git a/pkg/signatures/engine/engine.go b/pkg/signatures/engine/engine.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"sync"
 
+	"github.com/aquasecurity/tracee/pkg/events/findings"
 	"github.com/aquasecurity/tracee/pkg/logger"
 	"github.com/aquasecurity/tracee/pkg/signatures/metrics"
 	"github.com/aquasecurity/tracee/types/detect"
@@ -145,7 +146,12 @@ func (engine *Engine) unloadAllSignatures() {
 // matchHandler is a function that runs when a signature is matched
 func (engine *Engine) matchHandler(res *detect.Finding) {
 	_ = engine.stats.Detections.Increment()
+	e, err := findings.FindingToEvent(res)
+	if err != nil {
+		logger.Errorw("Failed to convert finding to event, will not feedback", "err", err)
+	}
 	engine.output <- res
+	engine.inputs.Tracee <- e.ToProtocol()
 }
 
 // checkCompletion is a function that runs at the end of each input source