mirror of
https://github.com/bettercap/bettercap
synced 2025-08-14 02:36:57 -07:00
new: graph module
This commit is contained in:
Simone Margaritelli
parent
6aa8f45d20
commit
db275429c2
7 changed files with 1078 additions and 0 deletions
399
modules/graph/graph.go
Normal file
399
modules/graph/graph.go
Normal file
|
|
@ -0,0 +1,399 @@
|
|||
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 Graph struct {
|
||||
sync.Mutex
|
||||
|
||||
path string
|
||||
}
|
||||
|
||||
func NewGraph(path string) (*Graph, error) {
|
||||
g := &Graph{
|
||||
path: path,
|
||||
}
|
||||
return g, nil
|
||||
}
|
||||
|
||||
func (g *Graph) EachNode(cb NodeCallback) error {
|
||||
g.Lock()
|
||||
defer g.Unlock()
|
||||
|
||||
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)
|
||||
} else {
|
||||
cb(&node)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
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")
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (g *Graph) Traverse(root string, onNode NodeCallback, onEdge EdgeCallback) error {
|
||||
if root == "" {
|
||||
// traverse the entire graph
|
||||
if err := g.EachNode(onNode); err != nil {
|
||||
return err
|
||||
} else if err = g.EachEdge(onEdge); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
// start by a specific node
|
||||
roots, err := g.FindOtherTypes("", root)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
stack := NewStack()
|
||||
|
||||
for _, root := range roots {
|
||||
stack.Push(root)
|
||||
}
|
||||
|
||||
type edgeBucket struct {
|
||||
left *Node
|
||||
edge *Edge
|
||||
right *Node
|
||||
}
|
||||
|
||||
allEdges := make([]edgeBucket, 0)
|
||||
visited := make(map[string]bool)
|
||||
|
||||
for {
|
||||
if last := stack.Pop(); last == nil {
|
||||
break
|
||||
} else {
|
||||
node := last.(*Node)
|
||||
nodeID := node.String()
|
||||
if _, found := visited[nodeID]; found {
|
||||
continue
|
||||
} else {
|
||||
visited[nodeID] = true
|
||||
}
|
||||
|
||||
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,
|
||||
right: right,
|
||||
})
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, edge := range allEdges {
|
||||
onEdge(edge.left, edge.edge, edge.right)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (g *Graph) Dot(filter, layout, name string) (string, error) {
|
||||
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
|
||||
}
|
||||
|
||||
data += "\n"
|
||||
data += " overlap=false\n"
|
||||
data += "}"
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func (g *Graph) FindNode(t NodeType, id string) (*Node, error) {
|
||||
g.Lock()
|
||||
defer g.Unlock()
|
||||
|
||||
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 nil, nil
|
||||
}
|
||||
|
||||
func (g *Graph) FindOtherTypes(t NodeType, id string) ([]*Node, error) {
|
||||
g.Lock()
|
||||
defer g.Unlock()
|
||||
|
||||
var otherNodes []*Node
|
||||
|
||||
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)
|
||||
} else {
|
||||
otherNodes = append(otherNodes, &node)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return otherNodes, nil
|
||||
}
|
||||
|
||||
func (g *Graph) CreateNode(t NodeType, id string, entity interface{}, annotations string) (*Node, error) {
|
||||
g.Lock()
|
||||
defer g.Unlock()
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
session.I.Events.Add("graph.node.new", node)
|
||||
|
||||
return node, nil
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (g *Graph) CreateEdge(from, to *Node, edgeType EdgeType) (*Edge, error) {
|
||||
g.Lock()
|
||||
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 {
|
||||
edges = append(edges, edge)
|
||||
} else {
|
||||
edges = []*Edge{edge}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
session.I.Events.Add("graph.edge.new", EdgeEvent{
|
||||
Left: from,
|
||||
Edge: edge,
|
||||
Right: to,
|
||||
})
|
||||
|
||||
return edge, nil
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue