diff --git a/_example/example.js b/_example/example.js index a7c4ad7c..753fa540 100644 --- a/_example/example.js +++ b/_example/example.js @@ -4,7 +4,7 @@ require("telegram") var fakeESSID = random.String(16, 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'); var fakeBSSID = random.Mac() -function graph(who, where) { +function createGraph(who, where) { // generates a .dot file with the graph for this mac run('graph.to_dot ' + who); // uses graphviz to make a png of it @@ -14,7 +14,7 @@ function graph(who, where) { function onDeauthentication(event) { var data = event.data; - graph(data.address1, '/tmp/graph_deauth.png'); + createGraph(data.address1, '/tmp/graph_deauth.png'); var message = '🚨 Detected deauthentication frame:\n\n' + // 'Time: ' + event.time + "\n" + @@ -34,7 +34,7 @@ function onDeauthentication(event) { function onNewAP(event){ var ap = event.data; if(ap.hostname == fakeESSID) { - graph(ap.mac, '/tmp/graph_ap.png'); + createGraph(ap.mac, '/tmp/graph_ap.png'); var message = '🦠 Detected rogue AP:\n\n' + // 'Time: ' + event.time + "\n" + @@ -52,7 +52,7 @@ function onHandshake(event){ var data = event.data; var what = 'handshake'; - graph(data.station, '/tmp/graph_handshake.png'); + createGraph(data.station, '/tmp/graph_handshake.png'); if(data.pmkid != null) { what = "RSN PMKID"; @@ -77,8 +77,8 @@ function onHandshake(event){ function onNewNode(event) { var node = event.data; - if(node.type != 'ssid' && node.type != 'ble_server') { - graph(node.id, '/tmp/graph_node.png'); + if(node.type != 'ssid' && node.type != 'ble_server' && graph.IsConnected(node.type, node.id)) { + createGraph(node.id, '/tmp/graph_node.png'); var message = '🖥️ Detected previously unknown ' + node.type + ':\n\n' + 'Type: ' + node.type + "\n" + diff --git a/modules/graph/graph.go b/modules/graph/graph.go index eecd96d9..43269c3f 100644 --- a/modules/graph/graph.go +++ b/modules/graph/graph.go @@ -10,6 +10,8 @@ import ( "time" ) +var Loaded = (* Graph)(nil) + type NodeCallback func(*Node) type EdgeCallback func(*Node, []Edge, *Node) @@ -24,10 +26,11 @@ func NewGraph(path string) (*Graph, error) { if edges, err := LoadEdges(path); err != nil { return nil, err } else { - return &Graph{ + Loaded = &Graph{ path: path, edges: edges, - }, nil + } + return Loaded, nil } } @@ -147,6 +150,10 @@ func (g *Graph) Traverse(root string, onNode NodeCallback, onEdge EdgeCallback) return nil } +func (g *Graph) IsConnected(nodeType string, nodeID string) bool { + return g.edges.IsConnected(fmt.Sprintf("%s_%s", nodeType, nodeID)) +} + func (g *Graph) Dot(filter, layout, name string, disconnected bool) (string, int, int, error) { size := 0 discarded := 0 diff --git a/modules/graph/js_builtin.go b/modules/graph/js_builtin.go new file mode 100644 index 00000000..49252d47 --- /dev/null +++ b/modules/graph/js_builtin.go @@ -0,0 +1,15 @@ +package graph + +import ( + "github.com/bettercap/bettercap/log" +) + +type graphPackage struct{} + +func (g graphPackage) IsConnected(nodeType, nodeID string) bool { + if Loaded == nil { + log.Error("graph.IsConnected: graph not loaded") + return false + } + return Loaded.IsConnected(nodeType, nodeID) +} diff --git a/modules/graph/module.go b/modules/graph/module.go index 14c49acd..211d24ca 100644 --- a/modules/graph/module.go +++ b/modules/graph/module.go @@ -6,6 +6,7 @@ import ( "github.com/bettercap/bettercap/network" "github.com/bettercap/bettercap/session" "github.com/evilsocket/islazy/fs" + "github.com/evilsocket/islazy/plugin" "github.com/evilsocket/islazy/str" "os" "path/filepath" @@ -50,6 +51,10 @@ type Module struct { wLock sync.Mutex } +func init() { + plugin.Defines["graph"] = graphPackage{} +} + func NewModule(s *session.Session) *Module { mod := &Module{ SessionModule: session.NewSessionModule("graph", s), @@ -173,8 +178,15 @@ func (mod *Module) updateSettings() error { } } - if mod.db, err = NewGraph(mod.settings.path); err != nil { - return err + // only reload if needed + if mod.db != nil && mod.db.path != mod.settings.path { + mod.db = nil + } + + if mod.db == nil { + if mod.db, err = NewGraph(mod.settings.path); err != nil { + return err + } } return nil diff --git a/modules/graph/node.go b/modules/graph/node.go index 314e0fd4..9f60b1cd 100644 --- a/modules/graph/node.go +++ b/modules/graph/node.go @@ -5,6 +5,7 @@ import ( "fmt" "io/ioutil" "os" + "strings" "time" "unicode" ) @@ -153,7 +154,7 @@ func (n Node) Dot(isTarget bool) string { return fmt.Sprintf("\"%s\" [%s, label=\"%s\"];", n.String(), style, - n.Label()) + strings.ReplaceAll(n.Label(), "\"", "\\\"")) } func (n Node) ToMap() (map[string]interface{}, error) { @@ -166,4 +167,4 @@ func (n Node) ToMap() (map[string]interface{}, error) { } return m, nil -} \ No newline at end of file +} diff --git a/modules/graph/dot.go b/modules/graph/to_dot.go similarity index 100% rename from modules/graph/dot.go rename to modules/graph/to_dot.go diff --git a/modules/graph/json.go b/modules/graph/to_json.go similarity index 100% rename from modules/graph/json.go rename to modules/graph/to_json.go diff --git a/session/session.go b/session/session.go index dd470f15..39e3cef1 100644 --- a/session/session.go +++ b/session/session.go @@ -3,7 +3,6 @@ package session import ( "errors" "fmt" - "github.com/evilsocket/islazy/plugin" "net" "os" "regexp" @@ -25,6 +24,7 @@ import ( "github.com/evilsocket/islazy/fs" "github.com/evilsocket/islazy/log" "github.com/evilsocket/islazy/ops" + "github.com/evilsocket/islazy/plugin" "github.com/evilsocket/islazy/str" "github.com/evilsocket/islazy/tui" )