feat: integrate custom detection rules with Armour

agent.go

@@ -67,7 +67,7 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 		return err
 	}
 
-	apiclient := &ApiClient{Client: &http.Client{Timeout: 3 * time.Second}, APIURL: config.APIURL, DisableTelemetry: config.DisableTelemetry, EgressPolicy: config.EgressPolicy, OneTimeKey: config.OneTimeKey}
+	apiclient := &ApiClient{Client: &http.Client{Timeout: 3 * time.Second}, APIURL: config.APIURL, TelemetryURL: config.TelemetryURL, DisableTelemetry: config.DisableTelemetry, EgressPolicy: config.EgressPolicy, OneTimeKey: config.OneTimeKey}
 
 	config.OneTimeKey = ""
 	// TODO: pass in an iowriter/ use log library
@@ -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{
@@ -233,22 +236,27 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 				APIURL:           config.APIURL,
 				Repo:             config.Repo,
 				CorrelationID:    config.CorrelationId,
-				OneTimeKey:       config.OneTimeKey,
+				OneTimeKey:       apiclient.OneTimeKey,
 				DisableTelemetry: config.DisableTelemetry,
 			},
+			EnableCustomDetectionRules: IsCustomDetectionRulesEnabled(),
 		}
 
 		conf.Files = append(conf.Files, getProcFilesOfInterest()...)
 
 		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")
+			}
 		}
 	}
 
@@ -370,10 +378,13 @@ func addImplicitEndpoints(endpoints map[string][]Endpoint, disableTelemetry bool
 		}
 	}
 
-	stepsecurity := Endpoint{domainName: "agent.api.stepsecurity.io", port: 443} // Should be implicit based on user feedback
+	stepsecurity := Endpoint{domainName: "agent.api.stepsecurity.io", port: 443}             // Should be implicit based on user feedback
+	stepsecurityTelemetry := Endpoint{domainName: "prod.app-api.stepsecurity.io", port: 443} // Telemetry endpoint for sending DNS and net connections to StepSecurity
+
 	if !disableTelemetry {
 		// allowing only if disable_telemetry is set to false
 		normalEndpoints[stepsecurity.domainName] = append(normalEndpoints[stepsecurity.domainName], stepsecurity)
+		normalEndpoints[stepsecurityTelemetry.domainName] = append(normalEndpoints[stepsecurityTelemetry.domainName], stepsecurityTelemetry)
 	}
 
 	return normalEndpoints, wildcardEndpoints

apiclient.go

@@ -41,6 +41,7 @@ type NetworkConnection struct {
 type ApiClient struct {
 	Client           *http.Client
 	APIURL           string
+	TelemetryURL     string
 	DisableTelemetry bool
 	EgressPolicy     string
 	OneTimeKey       string
@@ -57,7 +58,7 @@ func (apiclient *ApiClient) sendDNSRecord(correlationId, repo, domainName, ipAdd
 		dnsRecord.ResolvedIPAddress = ipAddress
 		dnsRecord.TimeStamp = time.Now().UTC()
 
-		url := fmt.Sprintf("%s/github/%s/actions/jobs/%s/dns", apiclient.APIURL, repo, correlationId)
+		url := fmt.Sprintf("%s/github/%s/actions/jobs/%s/dns", apiclient.TelemetryURL, repo, correlationId)
 
 		return apiclient.sendApiRequest("POST", url, dnsRecord)
 	}
@@ -76,7 +77,7 @@ func (apiclient *ApiClient) sendNetConnection(correlationId, repo, ipAddress, po
 		networkConnection.TimeStamp = timestamp
 		networkConnection.Tool = tool
 
-		url := fmt.Sprintf("%s/github/%s/actions/jobs/%s/networkconnection", apiclient.APIURL, repo, correlationId)
+		url := fmt.Sprintf("%s/github/%s/actions/jobs/%s/networkconnection", apiclient.TelemetryURL, repo, correlationId)
 
 		return apiclient.sendApiRequest("POST", url, networkConnection)
 	}

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")
+}

config.go

@@ -16,6 +16,7 @@ type config struct {
 	RunId                    string
 	WorkingDirectory         string
 	APIURL                   string
+	TelemetryURL             string
 	OneTimeKey               string
 	Endpoints                map[string][]Endpoint
 	EgressPolicy             string
@@ -37,6 +38,7 @@ type configFile struct {
 	RunId                    string `json:"run_id"`
 	WorkingDirectory         string `json:"working_directory"`
 	APIURL                   string `json:"api_url"`
+	TelemetryURL             string `json:"telemetry_url"`
 	OneTimeKey               string `json:"one_time_key"`
 	AllowedEndpoints         string `json:"allowed_endpoints"`
 	EgressPolicy             string `json:"egress_policy"`
@@ -65,6 +67,10 @@ func (c *config) init(configFilePath string) error {
 	c.RunId = configFile.RunId
 	c.WorkingDirectory = configFile.WorkingDirectory
 	c.APIURL = configFile.APIURL
+	c.TelemetryURL = configFile.TelemetryURL
+	if c.TelemetryURL == "" {
+		c.TelemetryURL = c.APIURL
+	}
 	c.Endpoints = parseEndpoints(configFile.AllowedEndpoints)
 	c.EgressPolicy = configFile.EgressPolicy
 	c.DisableTelemetry = configFile.DisableTelemetry

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 {
@@ -90,6 +91,8 @@ func (eventHandler *EventHandler) handleFileEvent(event *Event) {
 	}
 
 	eventHandler.fileMutex.Unlock()
+
+	eventHandler.submitFileEvent(event)
 }
 
 func isSourceCodeFile(fileName string) bool {
@@ -123,6 +126,8 @@ func (eventHandler *EventHandler) handleProcessEvent(event *Event) {
 	} else {
 		eventHandler.procMutex.Unlock()
 	}
+
+	eventHandler.submitProcessEvent(event)
 }
 
 /*
@@ -199,6 +204,8 @@ func (eventHandler *EventHandler) handleNetworkEvent(event *Event) {
 	}
 
 	eventHandler.netMutex.Unlock()
+
+	eventHandler.submitNetworkEvent(event)
 }
 
 func (eventHandler *EventHandler) HandleEvent(event *Event) {
@@ -440,3 +447,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
+}