feat: integrate custom detection rules with Armour

- Added support for custom detection rules in the Armour integration.
- Introduced new functions to submit process, file, and network events to the detection manager.
- Updated the DNS proxy to submit DNS events when custom detection rules are enabled.
- Refactored the agent's DNS handling to ensure proper execution flow when Docker is uninstalled.

Author: rohan-stepsecurity

agent.go

@@ -167,15 +167,18 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 	WriteLog("\n")
 	WriteLog("updated resolved")
 
-	// Change DNS for docker, causes process in containers to use agent's DNS proxy
-	if err := dnsConfig.SetDockerDNSServer(cmd, dockerDaemonConfigPath, tempDir); err != nil {
-		WriteLog(fmt.Sprintf("Error setting DNS server for docker %v", err))
-		RevertChanges(iptables, nflog, cmd, resolvdConfigPath, dockerDaemonConfigPath, dnsConfig, sudo)
-		return err
-	}
+	// we uninstall docker using go routine, handle case where that routine finishes before we come here
+	if !config.DisableSudoAndContainers {
+		// Change DNS for docker, causes process in containers to use agent's DNS proxy
+		if err := dnsConfig.SetDockerDNSServer(cmd, dockerDaemonConfigPath, tempDir); err != nil {
+			WriteLog(fmt.Sprintf("Error setting DNS server for docker %v", err))
+			RevertChanges(iptables, nflog, cmd, resolvdConfigPath, dockerDaemonConfigPath, dnsConfig, sudo)
+			return err
+		}
 
-	WriteLog("\n")
-	WriteLog("set docker config\n")
+		WriteLog("\n")
+		WriteLog("set docker config\n")
+	}
 
 	if config.EgressPolicy == EgressPolicyAudit {
 		netMonitor := NetworkMonitor{
@@ -242,13 +245,17 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 
 		conf.Files = append(conf.Files, getFilesOfInterest()...)
 
-		mArmour := armour.NewArmour(ctx, conf)
-		err := mArmour.Attach()
+		err := InitArmour(ctx, conf)
 		if err != nil {
 			WriteLog("Armour attachment failed")
 		} else {
-			defer mArmour.Detach()
+			if GlobalArmour != nil {
+				defer GlobalArmour.Detach()
+			}
 			WriteLog("Armour attached")
+			if IsCustomDetectionRulesEnabled() {
+				WriteLog("[armour] Custom detection rules enabled")
+			}
 		}
 	}
 

armour_manager.go

@@ -0,0 +1,31 @@
+package main
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/step-security/armour/armour"
+)
+
+// NOTE: before usage, make sure to nil check
+var GlobalArmour *armour.Armour = nil
+
+func InitArmour(ctx context.Context, conf *armour.Config) error {
+
+	GlobalArmour = armour.NewArmour(ctx, conf)
+	err := GlobalArmour.Init()
+	if err != nil {
+		GlobalArmour = nil
+		return err
+	}
+
+	runnerWorkerPID, err := getRunnerWorkerPID()
+	if err != nil {
+		WriteLog(fmt.Sprintf("[armour] Error getting Runner.Worker PID: %v", err))
+		return nil
+	}
+	GlobalArmour.SetRunnerWorkerPID(runnerWorkerPID)
+	WriteLog(fmt.Sprintf("[armour] Runner.Worker PID: %d", runnerWorkerPID))
+
+	return nil
+}

common.go

@@ -91,3 +91,7 @@ func getProcMemFiles(pid uint64) []string {
 
 	return out
 }
+
+func getRunnerWorkerPID() (uint64, error) {
+	return pidOf("Runner.Worker")
+}

dnsproxy.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/miekg/dns"
 	"github.com/pkg/errors"
+	"github.com/step-security/armour/armour"
 )
 
 type DNSProxy struct {
@@ -239,6 +240,8 @@ func (proxy *DNSProxy) getIPByDomain(domain string) (string, error) {
 
 	go proxy.ApiClient.sendDNSRecord(proxy.CorrelationId, proxy.Repo, domain, answer.Data)
 
+	go proxy.submitDNSEvent(answer.Data)
+
 	return answer.Data, nil
 
 }
@@ -296,6 +299,23 @@ func (proxy *DNSProxy) processTypeA(q *dns.Question, requestMsg *dns.Msg) (*dns.
 	return &rr, nil
 }
 
+// submitDNSEvent submits a DNS event to the detection manager.
+func (proxy *DNSProxy) submitDNSEvent(dest string) {
+	if !IsCustomDetectionRulesEnabled() {
+		return
+	}
+	if GlobalArmour == nil {
+		return
+	}
+	dm := GlobalArmour.DetectionManager()
+	if dm == nil {
+		return
+	}
+	dm.SubmitNetwork(&armour.NetworkDetectionEvent{
+		Dest: dest,
+	})
+}
+
 func startDNSServer(dnsProxy *DNSProxy, server DNSServer, errc chan error) {
 	dns.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
 		switch r.Opcode {

eventhandler.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/client"
+	"github.com/step-security/armour/armour"
 )
 
 type EventHandler struct {
@@ -89,6 +90,8 @@ func (eventHandler *EventHandler) handleFileEvent(event *Event) {
 		}
 	}
 
+	eventHandler.submitFileEvent(event)
+
 	eventHandler.fileMutex.Unlock()
 }
 
@@ -123,6 +126,8 @@ func (eventHandler *EventHandler) handleProcessEvent(event *Event) {
 	} else {
 		eventHandler.procMutex.Unlock()
 	}
+
+	eventHandler.submitProcessEvent(event)
 }
 
 /*
@@ -199,6 +204,10 @@ func (eventHandler *EventHandler) handleNetworkEvent(event *Event) {
 	}
 
 	eventHandler.netMutex.Unlock()
+
+	eventHandler.submitDNSEvent(reverseLookUp)
+
+	eventHandler.submitNetworkEvent(event)
 }
 
 func (eventHandler *EventHandler) HandleEvent(event *Event) {
@@ -440,3 +449,68 @@ func isPrivateIPAddress(ipAddress string) bool {
 func isIPv6(ip string) bool {
 	return strings.Contains(ip, ":")
 }
+
+func (eventHandler *EventHandler) submitProcessEvent(event *Event) {
+	if !IsCustomDetectionRulesEnabled() {
+		return
+	}
+	if GlobalArmour == nil {
+		return
+	}
+	dm := GlobalArmour.DetectionManager()
+	if dm == nil {
+		return
+	}
+	dm.SubmitProcess(&armour.ProcessDetectionEvent{
+		Pid:       event.Pid,
+		PPid:      event.PPid,
+		Exe:       event.Exe,
+		Arguments: event.ProcessArguments,
+		Cwd:       event.Path,
+		Timestamp: event.Timestamp,
+	})
+}
+
+// submitFileEvent submits a file event to the detection manager.
+func (eventHandler *EventHandler) submitFileEvent(event *Event) {
+	if !IsCustomDetectionRulesEnabled() {
+		return
+	}
+	if GlobalArmour == nil {
+		return
+	}
+	dm := GlobalArmour.DetectionManager()
+	if dm == nil {
+		return
+	}
+	dm.SubmitFile(&armour.FileDetectionEvent{
+		Syscall:   event.Syscall,
+		FileName:  filepath.Base(event.FileName),
+		Path:      event.FileName,
+		Exe:       event.Exe,
+		Pid:       event.Pid,
+		PPid:      event.PPid,
+		Timestamp: event.Timestamp,
+	})
+}
+
+// submitNetworkEvent submits a network event to the detection manager.
+func (eventHandler *EventHandler) submitNetworkEvent(event *Event) {
+	if GlobalArmour == nil {
+		return
+	}
+	dm := GlobalArmour.DetectionManager()
+	if dm == nil {
+		return
+	}
+
+	dm.SubmitNetwork(&armour.NetworkDetectionEvent{
+		Pid:       event.Pid,
+		PPid:      event.PPid,
+		Exe:       event.Exe,
+		Dest:      event.IPAddress,
+		DestIP:    event.IPAddress,
+		DestPort:  event.Port,
+		Timestamp: event.Timestamp,
+	})
+}

global_feature_flags.go

@@ -11,8 +11,9 @@ const (
 )
 
 type GlobalFeatureFlags struct {
-	AgentType    string `json:"agent_type"`
-	EnableArmour bool   `json:"enable_armour"`
+	AgentType                  string `json:"agent_type"`
+	EnableArmour               bool   `json:"enable_armour"`
+	EnableCustomDetectionRules bool   `json:"enable_custom_detection_rules"`
 }
 
 // GlobalFeatureFlagManager manages fetching and caching of global feature flags.
@@ -79,3 +80,8 @@ func IsArmourEnabled() bool {
 	flags := GetGlobalFeatureFlags()
 	return flags.EnableArmour
 }
+
+func IsCustomDetectionRulesEnabled() bool {
+	flags := GetGlobalFeatureFlags()
+	return flags.EnableCustomDetectionRules
+}