Merge pull request #97 from step-security/int

Add annotations for blocked traffic

Author: varunsh-coder

agent.go

@@ -5,8 +5,6 @@ import (
 	"fmt"
 	"net/http"
 	"os"
-	"sync"
-	"time"
 
 	"github.com/florianl/go-nflog/v2"
 )
@@ -47,8 +45,6 @@ type IPTables interface {
 	ClearChain(table, chain string) error
 }
 
-var fileMutex sync.Mutex
-
 // Run the agent
 // TODO: move all inputs into a struct
 func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
@@ -67,15 +63,15 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 	apiclient := &ApiClient{Client: &http.Client{}, APIURL: config.APIURL}
 
 	// TODO: pass in an iowriter/ use log library
-	writeLog(fmt.Sprintf("read config %v", config))
+	WriteLog(fmt.Sprintf("read config %v", config))
 
-	writeLog(fmt.Sprintf("%s %s", StepSecurityLogCorrelationPrefix, config.CorrelationId))
+	WriteLog(fmt.Sprintf("%s %s", StepSecurityLogCorrelationPrefix, config.CorrelationId))
 
 	// TODO: fix the cache and time
 	Cache := InitCache(10 * 60 * 1000000000) // 10 * 60 seconds
 
 	allowedEndpoints := addImplicitEndpoints(config.Endpoints)
-	
+
 	// Start DNS servers and get confirmation
 	dnsProxy := DNSProxy{
 		Cache:            &Cache,
@@ -93,20 +89,20 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 	if cmd == nil {
 		procMon := &ProcessMonitor{CorrelationId: config.CorrelationId, Repo: config.Repo, ApiClient: apiclient, WorkingDirectory: config.WorkingDirectory}
 		go procMon.MonitorProcesses(errc)
-		writeLog("started process monitor")
+		WriteLog("started process monitor")
 	}
 
 	dnsConfig := DnsConfig{}
 
 	var ipAddressEndpoints []ipAddressEndpoint
-	
+
 	// hydrate dns cache
 	if config.EgressPolicy == EgressPolicyBlock {
 		for _, endpoint := range allowedEndpoints {
 			// this will cause domain, IP mapping to be cached
 			ipAddress, err := dnsProxy.getIPByDomain(endpoint.domainName)
 			if err != nil {
-				writeLog(fmt.Sprintf("Error resolving allowed domain %v", err))
+				WriteLog(fmt.Sprintf("Error resolving allowed domain %v", err))
 				RevertChanges(iptables, nflog, cmd, resolvdConfigPath, dockerDaemonConfigPath, dnsConfig)
 				return err
 			}
@@ -118,21 +114,21 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 
 	// Change DNS config on host, causes processes to use agent's DNS proxy
 	if err := dnsConfig.SetDNSServer(cmd, resolvdConfigPath, tempDir); err != nil {
-		writeLog(fmt.Sprintf("Error setting DNS server %v", err))
+		WriteLog(fmt.Sprintf("Error setting DNS server %v", err))
 		RevertChanges(iptables, nflog, cmd, resolvdConfigPath, dockerDaemonConfigPath, dnsConfig)
 		return err
 	}
 
-	writeLog("updated resolved")
+	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))
+		WriteLog(fmt.Sprintf("Error setting DNS server for docker %v", err))
 		RevertChanges(iptables, nflog, cmd, resolvdConfigPath, dockerDaemonConfigPath, dnsConfig)
 		return err
 	}
 
-	writeLog("set docker config")
+	WriteLog("set docker config")
 
 	if config.EgressPolicy == EgressPolicyAudit {
 		netMonitor := NetworkMonitor{
@@ -144,20 +140,20 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 
 		// Start network monitor
 		go netMonitor.MonitorNetwork(nflog, errc) // listens for NFLOG messages
-		//writeLog("started net monitor")
-		writeLog("before audit rules")
+		
+		WriteLog("before audit rules")
 
 		// Add logging to firewall, including NFLOG rules
 		if err := AddAuditRules(iptables); err != nil {
-			writeLog(fmt.Sprintf("Error adding firewall rules %v", err))
+			WriteLog(fmt.Sprintf("Error adding firewall rules %v", err))
 			RevertChanges(iptables, nflog, cmd, resolvdConfigPath, dockerDaemonConfigPath, dnsConfig)
 			return err
 		}
 
-		writeLog("added audit rules")
+		WriteLog("added audit rules")
 	} else if config.EgressPolicy == EgressPolicyBlock {
 
-		writeLog(fmt.Sprintf("Allowed domains:%v", config.Endpoints))
+		WriteLog(fmt.Sprintf("Allowed domains:%v", config.Endpoints))
 
 		netMonitor := NetworkMonitor{
 			CorrelationId: config.CorrelationId,
@@ -170,13 +166,13 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 		go netMonitor.MonitorNetwork(nflog, errc) // listens for NFLOG messages
 
 		if err := addBlockRulesForGitHubHostedRunner(ipAddressEndpoints); err != nil {
-			writeLog(fmt.Sprintf("Error setting firewall for allowed domains %v", err))
+			WriteLog(fmt.Sprintf("Error setting firewall for allowed domains %v", err))
 			RevertChanges(iptables, nflog, cmd, resolvdConfigPath, dockerDaemonConfigPath, dnsConfig)
 			return err
 		}
 	}
 
-	writeLog("done")
+	WriteLog("done")
 
 	// Write the status file
 	writeStatus("Initialized")
@@ -186,7 +182,7 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 		case <-ctx.Done():
 			return nil
 		case e := <-errc:
-			writeLog(fmt.Sprintf("Error in Initialization %v", e))
+			WriteLog(fmt.Sprintf("Error in Initialization %v", e))
 			RevertChanges(iptables, nflog, cmd, resolvdConfigPath, dockerDaemonConfigPath, dnsConfig)
 			return e
 
@@ -210,29 +206,17 @@ func RevertChanges(iptables *Firewall, nflog AgentNflogger,
 	cmd Command, resolvdConfigPath, dockerDaemonConfigPath string, dnsConfig DnsConfig) {
 	err := RevertFirewallChanges(iptables)
 	if err != nil {
-		writeLog(fmt.Sprintf("Error in RevertChanges %v", err))
+		WriteLog(fmt.Sprintf("Error in RevertChanges %v", err))
 	}
 	err = dnsConfig.RevertDNSServer(cmd, resolvdConfigPath)
 	if err != nil {
-		writeLog(fmt.Sprintf("Error in reverting DNS server changes %v", err))
+		WriteLog(fmt.Sprintf("Error in reverting DNS server changes %v", err))
 	}
 	err = dnsConfig.RevertDockerDNSServer(cmd, dockerDaemonConfigPath)
 	if err != nil {
-		writeLog(fmt.Sprintf("Error in reverting docker DNS server changes %v", err))
+		WriteLog(fmt.Sprintf("Error in reverting docker DNS server changes %v", err))
 	}
-	writeLog("Reverted changes")
-}
-
-func writeLog(message string) {
-	fileMutex.Lock()
-	defer fileMutex.Unlock()
-
-	f, _ := os.OpenFile("/home/agent/agent.log",
-		os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
-
-	defer f.Close()
-
-	f.WriteString(fmt.Sprintf("%s:%s\n", time.Now().String(), message))
+	WriteLog("Reverted changes")
 }
 
 func writeStatus(message string) {

annotation.go

@@ -0,0 +1,21 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"sync"
+)
+
+var annotationMutex sync.Mutex
+
+func WriteAnnotation(message string) {
+	annotationMutex.Lock()
+	defer annotationMutex.Unlock()
+
+	f, _ := os.OpenFile("/home/agent/annotation.log",
+		os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+
+	defer f.Close()
+
+	f.WriteString(fmt.Sprintf("%s\n", message))
+}

dnsproxy.go

@@ -96,7 +96,8 @@ func (proxy *DNSProxy) getIPByDomain(domain string) (string, error) {
 
 	if proxy.EgressPolicy == EgressPolicyBlock {
 		if !proxy.isAllowedDomain(domain) {
-			go writeLog(fmt.Sprintf("domain not allowed: %s", domain))
+			go WriteLog(fmt.Sprintf("domain not allowed: %s", domain))
+			go WriteAnnotation(fmt.Sprintf("DNS resolution for domain %s was blocked", domain))
 			return "", fmt.Errorf("domain not allowed %s", domain)
 		}
 	}
@@ -127,7 +128,7 @@ func (proxy *DNSProxy) getIPByDomain(domain string) (string, error) {
 		if answer.Type == 1 {
 			proxy.Cache.Set(domain, answer.Data)
 
-			go writeLog(fmt.Sprintf("domain resolved: %s, ip address: %s", domain, answer.Data))
+			go WriteLog(fmt.Sprintf("domain resolved: %s, ip address: %s", domain, answer.Data))
 
 			go proxy.ApiClient.sendDNSRecord(proxy.CorrelationId, proxy.Repo, domain, answer.Data)
 

eventhandler.go

@@ -39,7 +39,7 @@ func (eventHandler *EventHandler) handleFileEvent(event *Event) {
 	}
 
 	if strings.Contains(event.FileName, "post_event.json") {
-		writeLog("post_event called")
+		WriteLog("post_event called")
 
 		// send done signal to post step
 		writeDone()
@@ -91,7 +91,6 @@ func (eventHandler *EventHandler) handleNetworkEvent(event *Event) {
 		_, found := eventHandler.ProcessConnectionMap[cacheKey]
 
 		if !found {
-			//writeLog(fmt.Sprintf("handleNetworkEvent %v", event))
 			tool := Tool{}
 			image := GetContainerByPid(event.Pid)
 			if image == "" {

log.go

@@ -0,0 +1,22 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"sync"
+	"time"
+)
+
+var logMutex sync.Mutex
+
+func WriteLog(message string) {
+	logMutex.Lock()
+	defer logMutex.Unlock()
+
+	f, _ := os.OpenFile("/home/agent/agent.log",
+		os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+
+	defer f.Close()
+
+	f.WriteString(fmt.Sprintf("%s:%s\n", time.Now().String(), message))
+}

main.go

@@ -34,12 +34,12 @@ func main() {
 				case syscall.SIGHUP:
 					c.init(agentConfigFilePath)
 				case os.Interrupt:
-					writeLog("got os.kill")
+					WriteLog("got os.kill")
 					cancel()
 					os.Exit(1)
 				}
 			case <-ctx.Done():
-				writeLog("called ctx.Done()")
+				WriteLog("called ctx.Done()")
 				os.Exit(1)
 			}
 		}

netmon.go

@@ -97,7 +97,8 @@ func (netMonitor *NetworkMonitor) handlePacket(attrs nflog.Attribute) {
 					ipv4.DstIP.String(), port, netMonitor.Status, timestamp, Tool{Name: Unknown, SHA256: Unknown})
 
 				if netMonitor.Status == "Dropped" {
-					go writeLog(fmt.Sprintf("ip address dropped: %s", ipv4.DstIP.String()))
+					go WriteLog(fmt.Sprintf("ip address dropped: %s", ipv4.DstIP.String()))
+					go WriteAnnotation(fmt.Sprintf("Traffic to IP Address %s was blocked", ipv4.DstIP.String()))
 				}
 			}
 		}

procmon.go

@@ -53,7 +53,7 @@ func (p *ProcessMonitor) PrepareEvent(sequence int, eventMap map[string]interfac
 	if !found {
 		p.Events[sequence] = &Event{}
 	}
-	//writeLog(fmt.Sprintf("%v", eventMap))
+
 	// Tags
 	value, found := eventMap["tags"]
 
@@ -94,7 +94,7 @@ func (p *ProcessMonitor) PrepareEvent(sequence int, eventMap map[string]interfac
 		argCountStr := fmt.Sprintf("%v", argc)
 		argCount, err := strconv.Atoi(argCountStr)
 		if err != nil {
-			writeLog(fmt.Sprintf("could not parse argc:%v", argc))
+			WriteLog(fmt.Sprintf("could not parse argc:%v", argc))
 		}
 		for i := 0; i < argCount; i++ {
 			p.Events[sequence].ProcessArguments = append(p.Events[sequence].ProcessArguments, fmt.Sprintf("%v", eventMap[fmt.Sprintf("a%d", i)]))

procmon_darwin.go

@@ -8,7 +8,7 @@ import (
 )
 
 func (p *ProcessMonitor) MonitorProcesses(errc chan error) {
-	writeLog("Monitor Processes called")
+	WriteLog("Monitor Processes called")
 }
 
 func getParentProcessId(pid string) (int, error) {