diff --git a/modules/graph/edge.go b/modules/graph/edge.go index d3466d41..cadd69a3 100644 --- a/modules/graph/edge.go +++ b/modules/graph/edge.go @@ -16,25 +16,26 @@ const ( ) type EdgeEvent struct { - Left *Node - Edge *Edge + Left *Node + Edge *Edge Right *Node } type Edge struct { - Type EdgeType `json:"type"` - CreatedAt time.Time `json:"created_at"` - Position session.GPS `json:"position"` + Type EdgeType `json:"type"` + CreatedAt time.Time `json:"created_at"` + Position *session.GPS `json:"position,omitempty"` } -func (e Edge) Dot(left, right *Node) string { - edgeLen := 2.0 +func (e Edge) Dot(left, right *Node, width float64) string { + 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) } diff --git a/modules/graph/edges.go b/modules/graph/edges.go new file mode 100644 index 00000000..eae8ac02 --- /dev/null +++ b/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() +} diff --git a/modules/graph/graph.go b/modules/graph/graph.go index 9e15ee2b..c49b5fcf 100644 --- a/modules/graph/graph.go +++ b/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{ - path: path, + if edges, err := LoadEdges(path); err != nil { + 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 raw, err := ioutil.ReadFile(fileName); err != nil { - 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) + if node, err := ReadNode(fileName); err != nil { + return 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 { - matches := edgesParser.FindAllStringSubmatch(path.Base(fileName), -1) - if len(matches) > 0 && len(matches[0]) == 3 { - var left, right Node - leftFileName := path.Join(g.path, matches[0][1]+".json") - rightFileName := path.Join(g.path, matches[0][2]+".json") + return g.edges.ForEachEdge(func(fromID string, edges []Edge, toID string) error { + var left, right *Node + var err error - if raw, err := ioutil.ReadFile(leftFileName); err != nil { - return fmt.Errorf("error while reading %s: %v", leftFileName, err) - } 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) - } + leftFileName := path.Join(g.path, fromID+".json") + rightFileName := path.Join(g.path, toID+".json") - var edges []*Edge - if raw, err := ioutil.ReadFile(fileName); err != nil { - return fmt.Errorf("error while reading %s: %v", fileName, err) - } else if err = json.Unmarshal(raw, &edges); err != nil { - 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) + if left, err = ReadNode(leftFileName); err != nil { + return err + } else if right, err = ReadNode(rightFileName); err != nil { + return err } + + 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 - edge *Edge + left *Node + 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 - edgesFilter := fmt.Sprintf("edges_%s_*.json", nodeID) - err = fs.Glob(g.path, edgesFilter, func(edgeFileName string) error { - right := new(Node) - - base := path.Base(edgeFileName) - base = strings.ReplaceAll(base, "edges_", "") - base = strings.ReplaceAll(base, nodeID + "_", "") - - // read right node - rightFileName := path.Join(g.path, base) - 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) - } - - 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, + // collect all edges starting from this node + err = g.edges.ForEachEdgeFrom(nodeID, func(_ string, edges []Edge, toID string) error { + rightFileName := path.Join(g.path, toID+".json") + if right, err := ReadNode(rightFileName); err != nil { + return err + } else { + // collect new node + if _, found := visited[toID]; !found { + stack.Push(right) + } + // collect all edges, we'll emit this later + allEdges = append(allEdges, edgeBucket{ + left: node, + edges: edges, 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) { - data += fmt.Sprintf(" %s\n", node.Dot(filter == node.ID)) - }, func(left *Node, edge *Edge, right *Node) { - data += fmt.Sprintf(" %s\n", edge.Dot(left, right)) - }); err != nil { - return "", err + typeMap := make(map[NodeType]bool) + + type typeCount struct { + edge Edge + count int } + 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 - 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 ReadNode(nodeFileName) } 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 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) + if node, err := ReadNode(nodeFileName); err != nil { + return nil, 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 { - return nil, fmt.Errorf("error creating data for %s: %v", nodeFileName, err) - } else if err = ioutil.WriteFile(nodeFileName, raw, os.ModePerm); err != nil { - return nil, fmt.Errorf("error creating %s: %v", nodeFileName, err) + if err := CreateNode(nodeFileName, node); err != nil { + return nil, 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 { - return fmt.Errorf("error creating new data for %s: %v", nodeFileName, err) - } else if err = ioutil.WriteFile(nodeFileName, raw, os.ModePerm); err != nil { - return fmt.Errorf("error updating %s: %v", nodeFileName, err) + if err := UpdateNode(nodeFileName, node); err != nil { + return 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() - defer g.Unlock() - - if _, edges, err := g.findEdgesUnlocked(from, to); err != nil { - return nil, err - } else { - num := len(edges) - for i := range edges { - // loop backwards - idx := num - 1 - i - edge := edges[idx] - if edge.Type == edgeType { - return edge, nil - } + edges := g.edges.FindEdges(from.String(), to.String(), true) + num := len(edges) + for i := range edges { + // 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() - defer g.Unlock() - - if _, edges, err := g.findEdgesUnlocked(from, to); err != nil { - return nil, err - } else { - num := len(edges) - for i := range edges { - // loop backwards - 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 + edges := g.edges.FindEdges(from.String(), to.String(), true) + num := len(edges) + for i := range edges { + // loop backwards + idx := num - 1 - i + edge := edges[idx] + if edge.Type == edgeType { + if time.Since(edge.CreatedAt) >= staleTime { + return nil, 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() - defer g.Unlock() - - var edgesFileName string - var edges []*Edge - - edge := &Edge{ + edge := Edge{ Type: edgeType, CreatedAt: time.Now(), - Position: session.I.GPS, } - if edgesFileName, edges, _ = g.findEdgesUnlocked(from, to); edges != nil { - edges = append(edges, edge) - } else { - edges = []*Edge{edge} + if session.I.GPS.Updated.IsZero() == false { + edge.Position = &session.I.GPS } - if raw, err := json.Marshal(edges); err != nil { - return nil, fmt.Errorf("error creating data for %s: %v", edgesFileName, err) - } else if err = ioutil.WriteFile(edgesFileName, raw, os.ModePerm); err != nil { - return nil, fmt.Errorf("error writing %s: %v", edgesFileName, err) + if err := g.edges.Connect(from.String(), to.String(), edge); err != nil { + return nil, err } session.I.Events.Add("graph.edge.new", EdgeEvent{ - Left: from, - Edge: edge, + Left: from, + Edge: &edge, Right: to, }) - return edge, nil + return &edge, nil } diff --git a/modules/graph/module.go b/modules/graph/module.go index 120fe2d3..c316625d 100644 --- a/modules/graph/module.go +++ b/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 - layout string - name string - output string + path string + layout string + name 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 { - return err - } else if manages == nil { - if manages, err = mod.db.CreateEdge(mod.gw, mod.iface, Manages); err != nil { + if iface.HwAddress == gw.HwAddress { + if err = mod.connectAsSame(mod.gw, mod.iface); 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 { + } else { + if manages, err := mod.db.FindLastRecentEdgeOfType(mod.gw, mod.iface, Manages, edgeStaleTime); 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 } diff --git a/modules/graph/node.go b/modules/graph/node.go index db71eed1..ed648426 100644 --- a/modules/graph/node.go +++ b/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", - BLEServer: "shape=box, style=filled, color=dodgerblue3", - Endpoint: "shape=box, style=filled, color=azure2", - Gateway: "shape=diamond, style=filled, color=azure4", - Station: "shape=box, style=filled, color=gold", - AccessPoint: "shape=diamond, style=filled, color=goldenrod3", + SSID: "shape=circle style=filled color=lightgray fillcolor=lightgray fixedsize=true penwidth=0.5", + BLEServer: "shape=box style=filled color=dodgerblue3", + Endpoint: "shape=box style=filled color=azure2", + Gateway: "shape=diamond style=filled color=azure4", + Station: "shape=box style=filled color=gold", + 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.String()) + n.Label()) } diff --git a/modules/graph/stack.go b/modules/graph/stack.go index 5652e68b..cf22db2f 100644 --- a/modules/graph/stack.go +++ b/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 { - stk := new(stack) +func NewStack() *Stack { + stk := new(Stack) stk.lock = &sync.Mutex{} return stk }