styx/matcher/main.go

package matcher

import (
	"bufio"
	"context"
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sync"
	"time"

	"github.com/dgraph-io/dgo/v2"
	"github.com/dgraph-io/dgo/v2/protos/api"
	"github.com/google/uuid"
	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"gitlab.dcso.lolcat/LABS/styx/models"
)

var (
	_, b, _, _ = runtime.Caller(0)
	basepath   = filepath.Dir(b)
)

// Matcher is the structure for the matching logic within Styx.
type Matcher struct {
	Running     bool
	StopChan    chan bool
	StoppedChan chan bool
}

// Initialize initialises the matcher based on the given configuration.
func (m *Matcher) Initialize() bool {
	if !viper.GetBool("matcher.activated") {
		return false
	}

	logrus.Info("matching is activated")
	return true
}

// Stop gracefully stops the matching logic.
func (m *Matcher) Stop(wg *sync.WaitGroup) {
	if m.Running {
		m.StopChan = make(chan bool)
		close(m.StopChan)
		<-m.StopChan
		wg.Done()
		m.Running = false
	}
}

// Result is the result from the matching query. Probably going to change.
type Result struct {
	Result []models.Node `json:"Node,omiempty"`
}

func loadTargets(graphClient *dgo.Dgraph) error {
	path := basepath + "/data/"

	sliceDomain, err := ioutil.ReadDir(path)
	if err != nil {
		logrus.Warn("matcher#ReadDir#domains", err)
		return err
	}

	for _, file := range sliceDomain {
		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()
			mu := &api.Mutation{
				CommitNow: true,
			}

			pb, err := json.Marshal(matcher)
			if err != nil {
				logrus.Error(err)
				return err
			}

			mu.SetJson = pb

			_, err = graphClient.NewTxn().Mutate(ctx, mu)
			if err != nil {
				logrus.Error(err)
				return err
			}

		}
		if err := scanner.Err(); err != nil {
			logrus.Error(err)
			return err
		}

	}

	return nil
}

// Run runs the routine trying to find matches in the ingested data.
func (m *Matcher) Run(wg *sync.WaitGroup, graphClient *dgo.Dgraph) {
	if err := loadTargets(graphClient); err != nil {
		logrus.Error(err)
	}
	// Created nodes based on the IOCs
	// Upsert those nodes if the values are found
	if !m.Running {
		// m.StoppedChan = make(chan bool)
		// wg.Add(1)
		// for {
		// q := `query allofterms($a: string) {
		// Node(func: allofterms(full, $a)) {
		// uid
		// }}`

		// ctx := context.Background()
		// txn := graphClient.NewTxn()
		// defer txn.Discard(ctx)
		// res, err := txn.QueryWithVars(ctx, q, map[string]string{"$a": "code"})
		// if err != nil {
		// logrus.Warn(err)
		// }

		// n := Result{}
		// json.Unmarshal([]byte(res.Json), &n)
		// uuid := uuid.New().String()
		// t := time.Now()
		// rfc3339time := t.Format(time.RFC3339)
		// matcher := models.Match{
		// ID:        uuid,
		// Timestamp: rfc3339time,
		// Target:    "code",
		// Nodes:     []models.Node{},
		// Type:      "matcher",
		// }
		// if len(n.Result) != 0 {
		// // TODO: review time and id to be updated on new resulsts

		// for _, res := range n.Result {
		// if len(matcher.Nodes) == 0 {
		// fmt.Println("First node appending")
		// matcher.Nodes = append(matcher.Nodes, res)
		// continue
		// }

		// for _, node := range matcher.Nodes {
		// if res.UID != node.UID {
		// fmt.Println("not there, appending...")
		// matcher.Nodes = append(matcher.Nodes, res)

		// }
		// }
		// }

		// fmt.Println("matcher:", matcher)

		// query := fmt.Sprintf(`
		// query  {
		// node as var(func: eq(id, %s))
		// }
		// `, matcher.ID)

		// nquads := fmt.Sprintf(`uid(node) <nodes> \"%v\"`, matcher.Nodes)

		// mu := &api.Mutation{
		// Cond:      fmt.Sprintf(`@if(eq(id, %s))`, matcher.ID),
		// SetNquads: []byte(nquads),
		// }

		// req := &api.Request{
		// Query:     query,
		// Mutations: []*api.Mutation{mu},
		// CommitNow: true,
		// }

		// ret, err := graphClient.NewTxn().Do(ctx, req)
		// fmt.Println("#####", ret)
		// if err != nil {
		// logrus.Fatal(err)
		// }
		// }

		m.Running = true
		// }
	}

}

// RunDomainMatch looks for a target within the identified IOCs in /matcher/data.
// the function could be refactored with RunDomainFilters
func RunDomainMatch(domain string) bool {
	path := basepath + "/data/domains.txt"

	sliceDomain, err := ioutil.ReadDir(path)
	if err != nil {
		logrus.Warn("matcher#ReadDir#domains", err)
	}

	for _, file := range sliceDomain {
		f, err := os.OpenFile(path+file.Name(), 1, 0644)
		if err != nil {
			logrus.Warn("matcher#OpenFile#", err)
		}
		scanner := bufio.NewScanner(f)

		for scanner.Scan() {
			r, err := regexp.Compile(scanner.Text())
			if err != nil {
				logrus.Warn("matcher#Compile#", err)
			}
			if r.MatchString(domain) {
				return false
			}
		}
	}
	return true
}