84e4937f85
This new work implements the server and the loader in two different binaries allowing the code while updating the IOC list. It updates also the documentation to reflect the new changes.
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{},
|
|
Type: "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!")
|
|
|
|
}
|