github/bettercap
1
0
Fork 0
mirror of https://github.com/bettercap/bettercap.git synced 2025-03-12 04:36:03 -07:00
Code Issues Packages Projects Releases Wiki Activity

misc: several improvements to the graph module

Browse Source
This commit is contained in:
Simone Margaritelli 2021-04-08 18:41:30 +02:00
parent 5b8cb9a82c
commit 71634058a7
6 changed files with 468 additions and 263 deletions

21
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)
}

158
modules/graph/edges.go Normal file

@ -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()
}

349
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
}

107
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
}

78
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())
}

18
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
}