Generate Graphviz graphs of network communication

analyze/analyzer.go

@@ -39,6 +39,9 @@ func Analyze(source *gopacket.PacketSource) error {
 			}
 		}
 
+		// Register communication for graph
+		output.AddPkgToGraph(packet)
+
 		// Raise statistics
 		totalPackets += 1
 		if processed {

main.go

@@ -36,6 +36,9 @@ func main() {
 	// Extract found and requested files
 	output.StoreFiles()
 
+	// Create communication graph
+	output.CreateGraph()
+
 	// Show user analysis
 	analyze.PrintSummary()
 

output/flag.go

@@ -3,11 +3,12 @@ package output
 import "flag"
 
 var (
-	fullOutput *bool
+	fullOutput       *bool
 	printEmptyBlocks *bool
-	targetFiles *string
-	targetAllFiles *bool
-	targetOutput *string
+	targetFiles      *string
+	targetAllFiles   *bool
+	targetOutput     *string
+	graphOutput      *string
 )
 
 func RegisterFlags() {
@@ -16,6 +17,5 @@ func RegisterFlags() {
 	targetFiles = flag.String("extract-these", "", "Comma-separated list of files to extract.")
 	targetAllFiles = flag.Bool("extract-all", false, "Extract all files found.")
 	targetOutput = flag.String("extract-to", "./extracted", "Directory to store extracted files in.")
+	graphOutput = flag.String("create-graph", "", "Create a Graphviz graph out of collected communication")
 }
-
-

output/graph.go

@@ -0,0 +1,104 @@
+package output
+
+import (
+	"fmt"
+	"io/ioutil"
+
+	"crypto/sha256"
+
+	"github.com/google/gopacket"
+)
+
+var graphPkgs []GraphPkg
+
+// AddPkgToGraph adds the given packet as communication to the graph
+func AddPkgToGraph(pkg gopacket.Packet) {
+	// Only proceed if pkg contains a network layer
+	if pkg.NetworkLayer() == nil {
+		return
+	}
+
+	src := pkg.NetworkLayer().NetworkFlow().Src().String()
+	dst := pkg.NetworkLayer().NetworkFlow().Dst().String()
+
+	// Search for the given communication pair
+	for _, p := range graphPkgs {
+		if p.from == src && p.to == dst {
+			// Communication pair found, add protocol and finish
+			p.AddProtocol("nil")
+			return
+		}
+	}
+
+	// Communcation pair was not in graphPkgs, add to it
+	graphPkgs = append(graphPkgs, GraphPkg{
+		from:     src,
+		to:       dst,
+		protocol: []string{""},
+	})
+}
+
+// CreateGraph writes out a Graphviz digraph
+func CreateGraph() {
+	if *graphOutput == "" {
+		// No graph requested
+		return
+	}
+
+	// Start with the Graphviz-specific header
+	dot := fmt.Sprintf("# Compile with `neato -Tpng %s > %s.png`\n", *graphOutput, *graphOutput)
+	dot += "digraph pancap {\n\toverlap = false;\n"
+
+	// First, gather all nodes as-is and write them out
+	dot += nodedef(graphPkgs)
+
+	// Iterate over communication
+	for _, p := range graphPkgs {
+		dot += fmt.Sprintf("\tn%s->n%s\n", hash(p.from), hash(p.to))
+	}
+
+	// Close
+	dot += "}\n"
+
+	// Write out
+	ioutil.WriteFile(*graphOutput, []byte(dot), 0644)
+}
+
+// Creates a list of distinct nodes, Graphviz-compatible
+func nodedef(pkgs []GraphPkg) string {
+	output := ""
+	nodes := []string{}
+	for _, p := range graphPkgs {
+		// Check if src and dst are already present in nodes array
+		srcFound := false
+		dstFound := false
+		for _, n := range nodes {
+			if p.from == n {
+				srcFound = true
+			}
+			if p.to == n {
+				dstFound = true
+			}
+		}
+		if !srcFound {
+			// src not yet present, add to node list
+			nodes = append(nodes, p.from)
+		}
+		if !dstFound {
+			// dst not yet present, add to node list
+			nodes = append(nodes, p.to)
+		}
+	}
+
+	// Output Graphviz-compatible node definition
+	for _, n := range nodes {
+		// As the Graphviz charset for nodes is rather small, rely on hashes
+		output += fmt.Sprintf("\tn%s[label=\"%s\"]\n", hash(n), n)
+	}
+
+	return output
+}
+
+func hash(s string) string {
+	return fmt.Sprintf("%x", sha256.Sum256([]byte(s)))[:6]
+}

output/graphpkg.go

@@ -0,0 +1,19 @@
+package output
+
+// GraphPkg resembles a directed communication from one address to another
+// It wraps up required information to draw a graph of the communication, including spoken protocols.
+type GraphPkg struct {
+	from     string
+	to       string
+	protocol []string
+}
+
+// AddProtocol adds the given protocol to the list of protocols if not already present
+func (p *GraphPkg) AddProtocol(protocol string) {
+	for _, p := range p.protocol {
+		if p == protocol {
+			return
+		}
+	}
+	p.protocol = append(p.protocol, protocol)
+}

output/output.go

@@ -23,4 +23,10 @@ func Finalize() {
 		// User avoided the files
 		printer.Println("Files found in stream. Add --extract-all or --extract-these <list> to extract them.")
 	}
+
+	// Check if something graph-worthy was collected
+	if *graphOutput == "" && len(graphPkgs) > 0 {
+		// User didn't want a graph
+		printer.Println("To summarize the communcation flow with a Graphviz graph, specify --create-graph <out.dot>.")
+	}
 }