new: exposing the graph object to js

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

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

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

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