118 lines
2.4 KiB
Go
118 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/json"
|
|
"io/ioutil"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"time"
|
|
|
|
"github.com/dgraph-io/dgo/v2"
|
|
"github.com/dgraph-io/dgo/v2/protos/api"
|
|
"github.com/google/uuid"
|
|
"github.com/sirupsen/logrus"
|
|
"gitlab.dcso.lolcat/LABS/styx/graph"
|
|
"gitlab.dcso.lolcat/LABS/styx/models"
|
|
)
|
|
|
|
var (
|
|
_, b, _, _ = runtime.Caller(0)
|
|
basepath = filepath.Dir(b)
|
|
)
|
|
|
|
// Result is the result from the matching query. Probably going to change.
|
|
type Result struct {
|
|
Result []models.Node `json:"Node,omitempty"`
|
|
}
|
|
|
|
func loadTargets(graphClient *dgo.Dgraph) error {
|
|
path := basepath + "/../../matcher/data/"
|
|
|
|
sliceDomain, err := ioutil.ReadDir(path)
|
|
if err != nil {
|
|
logrus.Warn("matcher#ReadDir#domains", err)
|
|
return err
|
|
}
|
|
|
|
for _, file := range sliceDomain {
|
|
logrus.Info("loading: ", file.Name(), " please wait...")
|
|
f, err := os.OpenFile(path+file.Name(), 0, 0644)
|
|
if err != nil {
|
|
logrus.Warn("matcher#OpenFile#", err)
|
|
return err
|
|
}
|
|
scanner := bufio.NewScanner(f)
|
|
|
|
for scanner.Scan() {
|
|
uuid := uuid.New().String()
|
|
t := time.Now()
|
|
rfc3339time := t.Format(time.RFC3339)
|
|
matcher := models.Match{
|
|
ID: uuid,
|
|
Timestamp: rfc3339time,
|
|
Target: scanner.Text(),
|
|
Nodes: []models.Node{},
|
|
NodeType: "matcher",
|
|
}
|
|
ctx := context.Background()
|
|
query := `query eq($a: string){
|
|
Node(func: eq(target, $a)){
|
|
uid
|
|
}
|
|
}`
|
|
|
|
txn := graphClient.NewTxn()
|
|
ret, err := txn.QueryWithVars(ctx, query, map[string]string{"$a": scanner.Text()})
|
|
if err != nil {
|
|
logrus.Warn(err)
|
|
}
|
|
|
|
n := Result{}
|
|
json.Unmarshal([]byte(ret.Json), &n)
|
|
|
|
// Check if the target already exists, if so, skipping not inserting
|
|
// the data
|
|
if len(n.Result) == 0 {
|
|
logrus.Info("new matcher, charging...")
|
|
mu := &api.Mutation{
|
|
CommitNow: true,
|
|
}
|
|
|
|
pb, err := json.Marshal(matcher)
|
|
if err != nil {
|
|
logrus.Error(err)
|
|
return err
|
|
}
|
|
|
|
mu.SetJson = pb
|
|
|
|
txn = graphClient.NewTxn()
|
|
defer txn.Discard(ctx)
|
|
|
|
_, err = txn.Mutate(ctx, mu)
|
|
if err != nil {
|
|
logrus.Error(err)
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func main() {
|
|
logrus.Info("Initializing Dgraph for the importer...")
|
|
dgraphClient, err := graph.ConnectToDgraph(false)
|
|
if err != nil || dgraphClient == nil {
|
|
logrus.WithField("err", err).Error("error initialising the graph database")
|
|
}
|
|
logrus.Info("Loading data...")
|
|
loadTargets(dgraphClient)
|
|
logrus.Info("Done!")
|
|
|
|
}
|