misc: several improvements to the graph module

modules/graph/edge.go

@@ -16,25 +16,26 @@ const (
 )
 
 type EdgeEvent struct {
-	Left *Node
+	Left  *Node
-	Edge *Edge
+	Edge  *Edge
 	Right *Node
 }
 
 type Edge struct {
-	Type      EdgeType    `json:"type"`
+	Type      EdgeType     `json:"type"`
-	CreatedAt time.Time   `json:"created_at"`
+	CreatedAt time.Time    `json:"created_at"`
-	Position  session.GPS `json:"position"`
+	Position  *session.GPS `json:"position,omitempty"`
 }
 
-func (e Edge) Dot(left, right *Node) string {
+func (e Edge) Dot(left, right *Node, width float64) string {
-	edgeLen := 2.0
+	edgeLen := 1.0
 	if e.Type == Is {
-		edgeLen = 1.0
+		edgeLen = 0.3
 	}
-	return fmt.Sprintf("\"%s\" -> \"%s\" [label=\"%s\", len=%.2f];",
+	return fmt.Sprintf("\"%s\" -> \"%s\" [label=\"%s\", len=%.2f, penwidth=%.2f];",
 		left.String(),
 		right.String(),
 		e.Type,
-		edgeLen)
+		edgeLen,
+		width)
 }

modules/graph/edges.go

@@ -0,0 +1,158 @@
+package graph
+
+import (
+	"encoding/json"
+	"github.com/evilsocket/islazy/fs"
+	"io/ioutil"
+	"os"
+	"path"
+	"sort"
+	"sync"
+	"time"
+)
+
+const edgesIndexName = "edges.json"
+
+type EdgesTo map[string][]Edge
+
+type EdgesCallback func(string, []Edge, string) error
+
+type Edges struct {
+	sync.RWMutex
+	timestamp time.Time
+	fileName  string
+	size      int
+	from      map[string]EdgesTo
+}
+
+type edgesJSON struct {
+	Timestamp time.Time          `json:"timestamp"`
+	Size      int                `json:"size"`
+	Edges     map[string]EdgesTo `json:"edges"`
+}
+
+func LoadEdges(basePath string) (*Edges, error) {
+	edges := Edges{
+		fileName: path.Join(basePath, edgesIndexName),
+		from:     make(map[string]EdgesTo),
+	}
+
+	if fs.Exists(edges.fileName) {
+		var js edgesJSON
+
+		if raw, err := ioutil.ReadFile(edges.fileName); err != nil {
+			return nil, err
+		} else if err = json.Unmarshal(raw, &js); err != nil {
+			return nil, err
+		}
+
+		edges.timestamp = js.Timestamp
+		edges.from = js.Edges
+		edges.size = js.Size
+	}
+
+	return &edges, nil
+}
+
+func (e *Edges) flush() error {
+	e.timestamp = time.Now()
+	js := edgesJSON{
+		Timestamp: e.timestamp,
+		Size:      e.size,
+		Edges:     e.from,
+	}
+
+	if raw, err := json.Marshal(js); err != nil {
+		return err
+	} else if err = ioutil.WriteFile(e.fileName, raw, os.ModePerm); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (e *Edges) Flush() error {
+	e.RLock()
+	defer e.RUnlock()
+	return e.flush()
+}
+
+func (e *Edges) ForEachEdge(cb EdgesCallback) error {
+	e.RLock()
+	defer e.RUnlock()
+
+	for from, edgesTo := range e.from {
+		for to, edges := range edgesTo {
+			if err := cb(from, edges, to); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+func (e *Edges) ForEachEdgeFrom(nodeID string, cb EdgesCallback) error {
+	e.RLock()
+	defer e.RUnlock()
+
+	if edgesTo, found := e.from[nodeID]; found {
+		for to, edges := range edgesTo {
+			if err := cb(nodeID, edges, to); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+func (e *Edges) IsConnected(nodeID string) bool {
+	e.RLock()
+	defer e.RUnlock()
+
+	if edgesTo, found := e.from[nodeID]; found {
+		return len(edgesTo) > 0
+	}
+
+	return false
+}
+
+func (e *Edges) FindEdges(fromID, toID string, doSort bool) []Edge {
+	e.RLock()
+	defer e.RUnlock()
+
+	if edgesTo, foundFrom := e.from[fromID]; foundFrom {
+		if edges, foundTo := edgesTo[toID]; foundTo {
+			if doSort {
+				// sort edges from oldest to newer
+				sort.Slice(edges, func(i, j int) bool {
+					return edges[i].CreatedAt.Before(edges[j].CreatedAt)
+				})
+			}
+			return edges
+		}
+	}
+
+	return nil
+}
+
+func (e *Edges) Connect(fromID, toID string, edge Edge) error {
+	e.Lock()
+	defer e.Unlock()
+
+	if edgesTo, foundFrom := e.from[fromID]; foundFrom {
+		edges := edgesTo[toID]
+		edges = append(edges, edge)
+		e.from[fromID][toID] = edges
+	} else {
+		// create the entire path
+		e.from[fromID] = EdgesTo{
+			toID: {edge},
+		}
+	}
+
+	e.size++
+
+	return e.flush()
+}

modules/graph/graph.go

@@ -1,36 +1,33 @@
 package graph
 
 import (
-	"encoding/json"
 	"fmt"
-	"strings"
 	"github.com/bettercap/bettercap/session"
 	"github.com/evilsocket/islazy/fs"
-	"io/ioutil"
-	"os"
 	"path"
-	"regexp"
-	"sort"
 	"sync"
 	"time"
 )
 
-var edgesParser = regexp.MustCompile(`^edges_(.+_[a-fA-F0-9:]{17})_(.+_.+)\.json$`)
-
 type NodeCallback func(*Node)
-type EdgeCallback func(*Node, *Edge, *Node)
+type EdgeCallback func(*Node, []Edge, *Node)
 
 type Graph struct {
 	sync.Mutex
 
-	path string
+	path  string
+	edges *Edges
 }
 
 func NewGraph(path string) (*Graph, error) {
-	g := &Graph{
+	if edges, err := LoadEdges(path); err != nil {
-		path: path,
+		return nil, err
+	} else {
+		return &Graph{
+			path:  path,
+			edges: edges,
+		}, nil
 	}
-	return g, nil
 }
 
 func (g *Graph) EachNode(cb NodeCallback) error {
@@ -39,13 +36,10 @@ func (g *Graph) EachNode(cb NodeCallback) error {
 
 	for _, nodeType := range NodeTypes {
 		err := fs.Glob(g.path, fmt.Sprintf("%s_*.json", nodeType), func(fileName string) error {
-			var node Node
+			if node, err := ReadNode(fileName); err != nil {
-			if raw, err := ioutil.ReadFile(fileName); err != nil {
+				return err
-				return fmt.Errorf("error while reading %s: %v", fileName, err)
-			} else if err = json.Unmarshal(raw, &node); err != nil {
-				return fmt.Errorf("error while decoding %s: %v", fileName, err)
 			} else {
-				cb(&node)
+				cb(node)
 			}
 			return nil
 		})
@@ -60,36 +54,21 @@ func (g *Graph) EachEdge(cb EdgeCallback) error {
 	g.Lock()
 	defer g.Unlock()
 
-	return fs.Glob(g.path, "edges_*.json", func(fileName string) error {
+	return g.edges.ForEachEdge(func(fromID string, edges []Edge, toID string) error {
-		matches := edgesParser.FindAllStringSubmatch(path.Base(fileName), -1)
+		var left, right *Node
-		if len(matches) > 0 && len(matches[0]) == 3 {
+		var err error
-			var left, right Node
-			leftFileName := path.Join(g.path, matches[0][1]+".json")
-			rightFileName := path.Join(g.path, matches[0][2]+".json")
 
-			if raw, err := ioutil.ReadFile(leftFileName); err != nil {
+		leftFileName := path.Join(g.path, fromID+".json")
-				return fmt.Errorf("error while reading %s: %v", leftFileName, err)
+		rightFileName := path.Join(g.path, toID+".json")
-			} else if err = json.Unmarshal(raw, &left); err != nil {
-				return fmt.Errorf("error while decoding %s: %v", leftFileName, err)
-			} else if raw, err = ioutil.ReadFile(rightFileName); err != nil {
-				return fmt.Errorf("error while reading %s: %v", rightFileName, err)
-			} else if err = json.Unmarshal(raw, &right); err != nil {
-				return fmt.Errorf("error while decoding %s: %v", rightFileName, err)
-			}
 
-			var edges []*Edge
+		if left, err = ReadNode(leftFileName); err != nil {
-			if raw, err := ioutil.ReadFile(fileName); err != nil {
+			return err
-				return fmt.Errorf("error while reading %s: %v", fileName, err)
+		} else if right, err = ReadNode(rightFileName); err != nil {
-			} else if err = json.Unmarshal(raw, &edges); err != nil {
+			return err
-				return fmt.Errorf("error while decoding %s: %v", fileName, err)
-			}
-
-			for _, edge := range edges {
-				cb(&left, edge, &right)
-			}
-		} else {
-			return fmt.Errorf("filename %s doesn't match edges parser", fileName)
 		}
+
+		cb(left, edges, right)
+
 		return nil
 	})
 }
@@ -110,14 +89,13 @@ func (g *Graph) Traverse(root string, onNode NodeCallback, onEdge EdgeCallback)
 		}
 
 		stack := NewStack()
-
 		for _, root := range roots {
 			stack.Push(root)
 		}
 
 		type edgeBucket struct {
-			left *Node
+			left  *Node
-			edge *Edge
+			edges []Edge
 			right *Node
 		}
 
@@ -138,74 +116,122 @@ func (g *Graph) Traverse(root string, onNode NodeCallback, onEdge EdgeCallback)
 
 				onNode(node)
 
-				// find all edges starting from this node
+				// collect all edges starting from this node
-				edgesFilter := fmt.Sprintf("edges_%s_*.json", nodeID)
+				err = g.edges.ForEachEdgeFrom(nodeID, func(_ string, edges []Edge, toID string) error {
-				err = fs.Glob(g.path, edgesFilter, func(edgeFileName string) error {
+					rightFileName := path.Join(g.path, toID+".json")
-					right := new(Node)
+					if right, err := ReadNode(rightFileName); err != nil {
-
+						return err
-					base := path.Base(edgeFileName)
+					} else {
-					base = strings.ReplaceAll(base, "edges_", "")
+						// collect new node
-					base = strings.ReplaceAll(base, nodeID + "_", "")
+						if _, found := visited[toID]; !found {
-
+							stack.Push(right)
-					// read right node
+						}
-					rightFileName := path.Join(g.path, base)
+						// collect all edges, we'll emit this later
-					if raw, err := ioutil.ReadFile(rightFileName); err != nil {
+						allEdges = append(allEdges, edgeBucket{
-						return fmt.Errorf("error while reading %s: %v", rightFileName, err)
+							left:  node,
-					} else if err = json.Unmarshal(raw, right); err != nil {
+							edges: edges,
-						return fmt.Errorf("error while decoding %s: %v", rightFileName, err)
-					}
-
-					stack.Push(right)
-
-					// read edges
-					var edges []*Edge
-					if raw, err := ioutil.ReadFile(edgeFileName); err != nil {
-						return fmt.Errorf("error while reading %s: %v", edgeFileName, err)
-					} else if err = json.Unmarshal(raw, &edges); err != nil {
-						return fmt.Errorf("error while decoding %s: %v", edgeFileName, err)
-					}
-
-					for _, edge := range edges {
-						allEdges = append(allEdges, edgeBucket {
-							left: node,
-							edge: edge,
 							right: right,
 						})
 					}
-
 					return nil
 				})
-				if err != nil {
-					return err
-				}
 			}
 		}
 
 		for _, edge := range allEdges {
-			onEdge(edge.left, edge.edge, edge.right)
+			onEdge(edge.left, edge.edges, edge.right)
 		}
 	}
 
 	return nil
 }
 
-func (g *Graph) Dot(filter, layout, name string) (string, error) {
+func (g *Graph) Dot(filter, layout, name string, disconnected bool) (string, int, int, error) {
+	size := 0
+	discarded := 0
+
 	data := fmt.Sprintf("digraph %s {\n", name)
 	data += fmt.Sprintf("  layout=%s\n", layout)
 
-	if err := g.Traverse(filter, func(node *Node) {
+	typeMap := make(map[NodeType]bool)
-		data += fmt.Sprintf("  %s\n", node.Dot(filter == node.ID))
+
-	}, func(left *Node, edge *Edge, right *Node) {
+	type typeCount struct {
-		data += fmt.Sprintf("  %s\n", edge.Dot(left, right))
+		edge  Edge
-	}); err != nil {
+		count int
-		return "", err
 	}
 
+	if err := g.Traverse(filter, func(node *Node) {
+		include := false
+		if disconnected || node.Type == SSID { // we don't create backwards edges for SSID
+			include = true
+		} else {
+			include = g.edges.IsConnected(node.String())
+		}
+
+		if include {
+			size++
+			typeMap[node.Type] = true
+			data += fmt.Sprintf("  %s\n", node.Dot(filter == node.ID))
+		} else {
+			discarded++
+		}
+	}, func(left *Node, edges []Edge, right *Node) {
+		// collect counters by edge type in order to calculate proportional widths
+		byType := make(map[string]typeCount)
+		tot := len(edges)
+
+		for _, edge := range edges {
+			if c, found := byType[string(edge.Type)]; found {
+				c.count++
+			} else {
+				byType[string(edge.Type)] = typeCount{
+					edge:  edge,
+					count: 1,
+				}
+			}
+		}
+
+		max := 2.0
+		for _, c := range byType {
+			w := max * float64(c.count/tot)
+			if w < 0.5 {
+				w = 0.5
+			}
+			data += fmt.Sprintf("  %s\n", c.edge.Dot(left, right, w))
+		}
+	}); err != nil {
+		return "", 0, 0, err
+	}
+
+	data += "\n"
+	data += "node [style=filled height=0.55 fontname=\"Verdana\" fontsize=10];\n"
+	data += "subgraph legend {\n" +
+			"graph[style=dotted];\n" +
+			"label = \"Legend\";\n"
+
+	var types []NodeType
+	for nodeType, _ := range typeMap {
+		types = append(types, nodeType)
+		node := Node{
+			Type:        nodeType,
+			Annotations: nodeTypeDescs[nodeType],
+			Dummy:       true,
+		}
+		data += fmt.Sprintf("  %s\n", node.Dot(false))
+	}
+
+	ntypes := len(types)
+	for i := 0; i < ntypes - 1; i++ {
+		data += fmt.Sprintf("  \"%s\" -> \"%s\" [style=invis];\n", types[i], types[i + 1])
+	}
+
+	data += "}\n"
+
 	data += "\n"
 	data += "  overlap=false\n"
 	data += "}"
 
-	return data, nil
+	return data, size, discarded, nil
 }
 
 func (g *Graph) FindNode(t NodeType, id string) (*Node, error) {
@@ -214,13 +240,7 @@ func (g *Graph) FindNode(t NodeType, id string) (*Node, error) {
 
 	nodeFileName := path.Join(g.path, fmt.Sprintf("%s_%s.json", t, id))
 	if fs.Exists(nodeFileName) {
-		var node Node
+		return ReadNode(nodeFileName)
-		if raw, err := ioutil.ReadFile(nodeFileName); err != nil {
-			return nil, fmt.Errorf("error while reading %s: %v", nodeFileName, err)
-		} else if err = json.Unmarshal(raw, &node); err != nil {
-			return nil, fmt.Errorf("error while decoding %s: %v", nodeFileName, err)
-		}
-		return &node, nil
 	}
 
 	return nil, nil
@@ -235,13 +255,10 @@ func (g *Graph) FindOtherTypes(t NodeType, id string) ([]*Node, error) {
 	for _, otherType := range NodeTypes {
 		if otherType != t {
 			if nodeFileName := path.Join(g.path, fmt.Sprintf("%s_%s.json", otherType, id)); fs.Exists(nodeFileName) {
-				var node Node
+				if node, err := ReadNode(nodeFileName); err != nil {
-				if raw, err := ioutil.ReadFile(nodeFileName); err != nil {
+					return nil, err
-					return nil, fmt.Errorf("error while reading %s: %v", nodeFileName, err)
-				} else if err = json.Unmarshal(raw, &node); err != nil {
-					return nil, fmt.Errorf("error while decoding %s: %v", nodeFileName, err)
 				} else {
-					otherNodes = append(otherNodes, &node)
+					otherNodes = append(otherNodes, node)
 				}
 			}
 		}
@@ -257,16 +274,13 @@ func (g *Graph) CreateNode(t NodeType, id string, entity interface{}, annotation
 	node := &Node{
 		Type:        t,
 		ID:          id,
-		CreatedAt:   time.Now(),
 		Entity:      entity,
 		Annotations: annotations,
 	}
 
 	nodeFileName := path.Join(g.path, fmt.Sprintf("%s.json", node.String()))
-	if raw, err := json.Marshal(node); err != nil {
+	if err := CreateNode(nodeFileName, node); err != nil {
-		return nil, fmt.Errorf("error creating data for %s: %v", nodeFileName, err)
+		return nil, err
-	} else if err = ioutil.WriteFile(nodeFileName, raw, os.ModePerm); err != nil {
-		return nil, fmt.Errorf("error creating %s: %v", nodeFileName, err)
 	}
 
 	session.I.Events.Add("graph.node.new", node)
@@ -278,86 +292,40 @@ func (g *Graph) UpdateNode(node *Node) error {
 	g.Lock()
 	defer g.Unlock()
 
-	node.UpdatedAt = time.Now()
 	nodeFileName := path.Join(g.path, fmt.Sprintf("%s.json", node.String()))
-	if raw, err := json.Marshal(node); err != nil {
+	if err := UpdateNode(nodeFileName, node); err != nil {
-		return fmt.Errorf("error creating new data for %s: %v", nodeFileName, err)
+		return err
-	} else if err = ioutil.WriteFile(nodeFileName, raw, os.ModePerm); err != nil {
-		return fmt.Errorf("error updating %s: %v", nodeFileName, err)
 	}
 
 	return nil
 }
 
-func (g *Graph) findEdgesUnlocked(from, to *Node) (string, []*Edge, error) {
-	edgesFileName := path.Join(g.path, fmt.Sprintf("edges_%s_%s.json", from.String(), to.String()))
-	if fs.Exists(edgesFileName) {
-		var edges []*Edge
-		if raw, err := ioutil.ReadFile(edgesFileName); err != nil {
-			return edgesFileName, nil, fmt.Errorf("error while reading %s: %v", edgesFileName, err)
-		} else if err = json.Unmarshal(raw, &edges); err != nil {
-			return edgesFileName, nil, fmt.Errorf("error while decoding %s: %v", edgesFileName, err)
-		}
-
-		// sort edges from oldest to newer
-		sort.Slice(edges, func(i, j int) bool {
-			return edges[i].CreatedAt.Before(edges[j].CreatedAt)
-		})
-
-		return edgesFileName, edges, nil
-	}
-
-	return edgesFileName, nil, nil
-}
-
-func (g *Graph) FindEdges(from, to *Node) ([]*Edge, error) {
-	g.Lock()
-	defer g.Unlock()
-
-	_, edges, err := g.findEdgesUnlocked(from, to)
-	return edges, err
-}
-
 func (g *Graph) FindLastEdgeOfType(from, to *Node, edgeType EdgeType) (*Edge, error) {
-	g.Lock()
+	edges := g.edges.FindEdges(from.String(), to.String(), true)
-	defer g.Unlock()
+	num := len(edges)
-
+	for i := range edges {
-	if _, edges, err := g.findEdgesUnlocked(from, to); err != nil {
+		// loop backwards
-		return nil, err
+		idx := num - 1 - i
-	} else {
+		edge := edges[idx]
-		num := len(edges)
+		if edge.Type == edgeType {
-		for i := range edges {
+			return &edge, nil
-			// loop backwards
-			idx := num - 1 - i
-			edge := edges[idx]
-			if edge.Type == edgeType {
-				return edge, nil
-			}
 		}
 	}
-
 	return nil, nil
 }
 
 func (g *Graph) FindLastRecentEdgeOfType(from, to *Node, edgeType EdgeType, staleTime time.Duration) (*Edge, error) {
-	g.Lock()
+	edges := g.edges.FindEdges(from.String(), to.String(), true)
-	defer g.Unlock()
+	num := len(edges)
-
+	for i := range edges {
-	if _, edges, err := g.findEdgesUnlocked(from, to); err != nil {
+		// loop backwards
-		return nil, err
+		idx := num - 1 - i
-	} else {
+		edge := edges[idx]
-		num := len(edges)
+		if edge.Type == edgeType {
-		for i := range edges {
+			if time.Since(edge.CreatedAt) >= staleTime {
-			// loop backwards
+				return nil, nil
-			idx := num - 1 - i
-			edge := edges[idx]
-			if edge.Type == edgeType {
-				if time.Since(edge.CreatedAt) >= staleTime {
-					return nil, nil
-				}
-				// edge is still fresh
-				return edge, nil
 			}
+			return &edge, nil
 		}
 	}
 
@@ -365,35 +333,24 @@ func (g *Graph) FindLastRecentEdgeOfType(from, to *Node, edgeType EdgeType, stal
 }
 
 func (g *Graph) CreateEdge(from, to *Node, edgeType EdgeType) (*Edge, error) {
-	g.Lock()
+	edge := Edge{
-	defer g.Unlock()
-
-	var edgesFileName string
-	var edges []*Edge
-
-	edge := &Edge{
 		Type:      edgeType,
 		CreatedAt: time.Now(),
-		Position:  session.I.GPS,
 	}
 
-	if edgesFileName, edges, _ = g.findEdgesUnlocked(from, to); edges != nil {
+	if session.I.GPS.Updated.IsZero() == false {
-		edges = append(edges, edge)
+		edge.Position = &session.I.GPS
-	} else {
-		edges = []*Edge{edge}
 	}
 
-	if raw, err := json.Marshal(edges); err != nil {
+	if err := g.edges.Connect(from.String(), to.String(), edge); err != nil {
-		return nil, fmt.Errorf("error creating data for %s: %v", edgesFileName, err)
+		return nil, err
-	} else if err = ioutil.WriteFile(edgesFileName, raw, os.ModePerm); err != nil {
-		return nil, fmt.Errorf("error writing %s: %v", edgesFileName, err)
 	}
 
 	session.I.Events.Add("graph.edge.new", EdgeEvent{
-		Left: from,
+		Left:  from,
-		Edge: edge,
+		Edge:  &edge,
 		Right: to,
 	})
 
-	return edge, nil
+	return &edge, nil
 }

modules/graph/module.go

@@ -11,6 +11,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"regexp"
 	"time"
 )
 
@@ -19,11 +20,15 @@ const (
 	edgeStaleTime   = time.Hour * 24
 )
 
+var privacyFilter = regexp.MustCompile("(?i)([a-f0-9]{2}):([a-f0-9]{2}):([a-f0-9]{2}):([a-f0-9]{2}):([a-f0-9]{2}):([a-f0-9]{2})")
+
 type settings struct {
-	path string
+	path         string
-	layout string
+	layout       string
-	name string
+	name         string
-	output string
+	output       string
+	disconnected bool
+	privacy      bool
 }
 
 type Module struct {
@@ -40,9 +45,9 @@ func NewModule(s *session.Session) *Module {
 	mod := &Module{
 		SessionModule: session.NewSessionModule("graph", s),
 		settings: settings{
-			path: filepath.Join(caplets.InstallBase, "graph"),
+			path:   filepath.Join(caplets.InstallBase, "graph"),
 			layout: "neato",
-			name: "bettergraph",
+			name:   "bettergraph",
 			output: "bettergraph.dot",
 		},
 	}
@@ -67,6 +72,14 @@ func NewModule(s *session.Session) *Module {
 		"",
 		"File name for dot output."))
 
+	mod.AddParam(session.NewBoolParameter("graph.dot.disconnected",
+		"false",
+		"Include disconnected edges in then output graph."))
+
+	mod.AddParam(session.NewBoolParameter("graph.dot.privacy",
+		"false",
+		"Obfuscate mac addresses."))
+
 	mod.AddHandler(session.NewModuleHandler("graph on", "",
 		"Start the Module module.",
 		func(args []string) error {
@@ -114,6 +127,10 @@ func (mod *Module) updateSettings() error {
 		return err
 	} else if err, mod.settings.output = mod.StringParam("graph.dot.output"); err != nil {
 		return err
+	} else if err, mod.settings.disconnected = mod.BoolParam("graph.dot.disconnected"); err != nil {
+		return err
+	} else if err, mod.settings.privacy = mod.BoolParam("graph.dot.privacy"); err != nil {
+		return err
 	} else if err, mod.settings.path = mod.StringParam("graph.path"); err != nil {
 		return err
 	} else if mod.settings.path, err = filepath.Abs(mod.settings.path); err != nil {
@@ -140,6 +157,19 @@ func (mod *Module) Configure() (err error) {
 
 	// if have an IP
 	if mod.Session.Gateway != nil && mod.Session.Interface != nil {
+		// find or create interface node
+		iface := mod.Session.Interface
+		if mod.iface, err = mod.db.FindNode(Endpoint, iface.HwAddress); err != nil {
+			return err
+		} else if mod.iface == nil {
+			// create the interface node
+			if mod.iface, err = mod.db.CreateNode(Endpoint, iface.HwAddress, iface, ifaceAnnotation); err != nil {
+				return err
+			}
+		} else if err = mod.db.UpdateNode(mod.iface); err != nil {
+			return err
+		}
+
 		// find or create gateway node
 		gw := mod.Session.Gateway
 		if mod.gw, err = mod.db.FindNode(Gateway, gw.HwAddress); err != nil {
@@ -154,33 +184,26 @@ func (mod *Module) Configure() (err error) {
 			}
 		}
 
-		// find or create interface node
-		iface := mod.Session.Interface
-		if mod.iface, err = mod.db.FindNode(Endpoint, iface.HwAddress); err != nil {
-			return err
-		} else if mod.iface == nil {
-			// create the interface node
-			if mod.iface, err = mod.db.CreateNode(Endpoint, iface.HwAddress, iface, ifaceAnnotation); err != nil {
-				return err
-			}
-		} else if err = mod.db.UpdateNode(mod.iface); err != nil {
-			return err
-		}
-
 		// create relations if needed
-		if manages, err := mod.db.FindLastRecentEdgeOfType(mod.gw, mod.iface, Manages, edgeStaleTime); err != nil {
+		if iface.HwAddress == gw.HwAddress {
-			return err
+			if err = mod.connectAsSame(mod.gw, mod.iface); err != nil {
-		} else if manages == nil {
-			if manages, err = mod.db.CreateEdge(mod.gw, mod.iface, Manages); err != nil {
 				return err
 			}
-		}
+		} else {
-
+			if manages, err := mod.db.FindLastRecentEdgeOfType(mod.gw, mod.iface, Manages, edgeStaleTime); err != nil {
-		if connects_to, err := mod.db.FindLastEdgeOfType(mod.iface, mod.gw, ConnectsTo); err != nil {
-			return err
-		} else if connects_to == nil {
-			if connects_to, err = mod.db.CreateEdge(mod.iface, mod.gw, ConnectsTo); err != nil {
 				return err
+			} else if manages == nil {
+				if manages, err = mod.db.CreateEdge(mod.gw, mod.iface, Manages); err != nil {
+					return err
+				}
+			}
+
+			if connects_to, err := mod.db.FindLastEdgeOfType(mod.iface, mod.gw, ConnectsTo); err != nil {
+				return err
+			} else if connects_to == nil {
+				if connects_to, err = mod.db.CreateEdge(mod.iface, mod.gw, ConnectsTo); err != nil {
+					return err
+				}
 			}
 		}
 	}
@@ -192,16 +215,28 @@ func (mod *Module) Configure() (err error) {
 
 func (mod *Module) generateDotGraph(bssid string) error {
 	start := time.Now()
-
 	if err := mod.updateSettings(); err != nil {
 		return err
-	} else if data, err := mod.db.Dot(bssid, mod.settings.layout, mod.settings.name); err != nil {
-		return err
-	} else if err := ioutil.WriteFile(mod.settings.output, []byte(data), os.ModePerm); err != nil {
-		return err
 	}
 
-	mod.Info("graph saved to %s in %v", mod.settings.output, time.Since(start))
+	data, size, discarded, err := mod.db.Dot(bssid, mod.settings.layout, mod.settings.name, mod.settings.disconnected)
+	if err != nil {
+		return err
+	}
+
+	if mod.settings.privacy {
+		data = privacyFilter.ReplaceAllString(data, "$1:$2:xx:xx:xx:xx")
+	}
+
+	if err := ioutil.WriteFile(mod.settings.output, []byte(data), os.ModePerm); err != nil {
+		return err
+	} else {
+		mod.Info("graph saved to %s in %v (%d edges, %d discarded)",
+			mod.settings.output,
+			time.Since(start),
+			size,
+			discarded)
+	}
 	return nil
 }
 
@@ -317,7 +352,7 @@ func (mod *Module) createDot11Graph(ap *network.AccessPoint, station *network.St
 }
 
 func (mod *Module) createDot11ProbeGraph(ssid string, station *network.Station) (*Node, bool, *Node, bool, error) {
-	apNode, apIsNew, err := mod.createDot11SSIDGraph(station.HwAddress + fmt.Sprintf(":PROBE:%x", ssid), ssid)
+	apNode, apIsNew, err := mod.createDot11SSIDGraph(station.HwAddress+fmt.Sprintf(":PROBE:%x", ssid), ssid)
 	if err != nil {
 		return nil, false, nil, false, err
 	}

modules/graph/node.go

@@ -1,7 +1,10 @@
 package graph
 
 import (
+	"encoding/json"
 	"fmt"
+	"io/ioutil"
+	"os"
 	"time"
 	"unicode"
 )
@@ -26,13 +29,22 @@ var NodeTypes = []NodeType{
 	BLEServer,
 }
 
+var nodeTypeDescs = map[NodeType]string{
+	SSID:        "WiFI SSID probe",
+	BLEServer:   "BLE Device",
+	Station:     "WiFi Client",
+	AccessPoint: "WiFi AP",
+	Endpoint:    "IP Client",
+	Gateway:     "IP Gateway",
+}
+
 var nodeDotStyles = map[NodeType]string{
-	SSID:        "shape=diamond",
+	SSID:        "shape=circle style=filled color=lightgray fillcolor=lightgray fixedsize=true penwidth=0.5",
-	BLEServer:   "shape=box, style=filled, color=dodgerblue3",
+	BLEServer:   "shape=box style=filled color=dodgerblue3",
-	Endpoint:    "shape=box, style=filled, color=azure2",
+	Endpoint:    "shape=box style=filled color=azure2",
-	Gateway:     "shape=diamond, style=filled, color=azure4",
+	Gateway:     "shape=diamond style=filled color=azure4",
-	Station:     "shape=box, style=filled, color=gold",
+	Station:     "shape=box style=filled color=gold",
-	AccessPoint: "shape=diamond, style=filled, color=goldenrod3",
+	AccessPoint: "shape=diamond style=filled color=goldenrod3",
 }
 
 type Node struct {
@@ -42,13 +54,54 @@ type Node struct {
 	ID          string      `json:"id"`
 	Annotations string      `json:"annotations"`
 	Entity      interface{} `json:"entity"`
+	Dummy       bool        `json:"-"`
+}
+
+func ReadNode(fileName string) (*Node, error) {
+	var node Node
+	if raw, err := ioutil.ReadFile(fileName); err != nil {
+		return nil, fmt.Errorf("error while reading %s: %v", fileName, err)
+	} else if err = json.Unmarshal(raw, &node); err != nil {
+		return nil, fmt.Errorf("error while decoding %s: %v", fileName, err)
+	}
+	return &node, nil
+}
+
+func WriteNode(fileName string, node *Node, update bool) error {
+	if update {
+		node.UpdatedAt = time.Now()
+	} else {
+		node.CreatedAt = time.Now()
+	}
+
+	if raw, err := json.Marshal(node); err != nil {
+		return fmt.Errorf("error creating data for %s: %v", fileName, err)
+	} else if err = ioutil.WriteFile(fileName, raw, os.ModePerm); err != nil {
+		return fmt.Errorf("error creating %s: %v", fileName, err)
+	}
+	return nil
+}
+
+func CreateNode(fileName string, node *Node) error {
+	return WriteNode(fileName, node, false)
+}
+
+func UpdateNode(fileName string, node *Node) error {
+	return WriteNode(fileName, node, true)
 }
 
 func (n Node) String() string {
-	return fmt.Sprintf("%s_%s", n.Type, n.ID)
+	if n.Dummy == false {
+		return fmt.Sprintf("%s_%s", n.Type, n.ID)
+	}
+	return string(n.Type)
 }
 
 func (n Node) Label() string {
+	if n.Dummy {
+		return n.Annotations
+	}
+
 	switch n.Type {
 	case SSID:
 		s := n.Entity.(string)
@@ -74,9 +127,10 @@ func (n Node) Label() string {
 			n.Entity.(map[string]interface{})["mac"].(string),
 			n.Entity.(map[string]interface{})["vendor"].(string))
 	case AccessPoint:
-		return fmt.Sprintf("%s\\n(%s)",
+		return fmt.Sprintf("%s\\n%s\\n(%s)",
 			n.Entity.(map[string]interface{})["hostname"].(string),
-			n.Entity.(map[string]interface{})["mac"].(string))
+			n.Entity.(map[string]interface{})["mac"].(string),
+			n.Entity.(map[string]interface{})["vendor"].(string))
 	case Endpoint:
 		return fmt.Sprintf("%s\\n(%s %s)",
 			n.Entity.(map[string]interface{})["ipv4"].(string),
@@ -96,8 +150,8 @@ func (n Node) Dot(isTarget bool) string {
 	if isTarget {
 		style += ", color=red"
 	}
-	return fmt.Sprintf("node [%s]; {node [label=\"%s\"] \"%s\";};",
+	return fmt.Sprintf("\"%s\" [%s, label=\"%s\"];",
+		n.String(),
 		style,
-		n.Label(),
+		n.Label())
-		n.String())
 }

modules/graph/stack.go

@@ -2,21 +2,21 @@ package graph
 
 import "sync"
 
-type element struct {
+type entry struct {
 	data interface{}
-	next *element
+	next *entry
 }
 
-type stack struct {
+type Stack struct {
 	lock *sync.Mutex
-	head *element
+	head *entry
 	Size int
 }
 
-func (stk *stack) Push(data interface{}) {
+func (stk *Stack) Push(data interface{}) {
 	stk.lock.Lock()
 
-	element := new(element)
+	element := new(entry)
 	element.data = data
 	temp := stk.head
 	element.next = temp
@@ -26,7 +26,7 @@ func (stk *stack) Push(data interface{}) {
 	stk.lock.Unlock()
 }
 
-func (stk *stack) Pop() interface{} {
+func (stk *Stack) Pop() interface{} {
 	if stk.head == nil {
 		return nil
 	}
@@ -40,8 +40,8 @@ func (stk *stack) Pop() interface{} {
 	return r
 }
 
-func NewStack() *stack {
+func NewStack() *Stack {
-	stk := new(stack)
+	stk := new(Stack)
 	stk.lock = &sync.Mutex{}
 	return stk
 }