bettercap/modules/graph/edges.go

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