package matcher

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"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,omitempty"`
}

func loadTargets(graphClient *dgo.Dgraph) ([]string, error) {
	path := basepath + "/data/"
	res := []string{}

	sliceDomain, err := ioutil.ReadDir(path)
	if err != nil {
		logrus.Warn("matcher#ReadDir#domains", err)
		return nil, 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 nil, 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 nil, err
				}

				mu.SetJson = pb

				txn = graphClient.NewTxn()
				defer txn.Discard(ctx)

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

			}

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

		}
	}

	return res, nil
}

// Run runs the routine trying to find matches in the ingested Pastebin data.
func (m *Matcher) Run(wg *sync.WaitGroup, graphClient *dgo.Dgraph) {
	logrus.Info("loading matcher targets, this might take some time...")
	targets, err := loadTargets(graphClient)
	if err != nil {
		logrus.Error(err)
	}
	logrus.Info("finished loading matcher targets")
	fmt.Println(targets)

	if !m.Running {
		m.StoppedChan = make(chan bool)
		wg.Add(1)
		logrus.Info("Running matchers")
		for _, target := range targets {
			go runPasteMatcher(target, graphClient)
			go runCertstreamMatcher(target, graphClient)
			go runShodanMatcher(target, graphClient)
		}

		// TODO: probably not the best design here
		wg.Add(len(targets))
		m.Running = true
	}
}

func runPasteMatcher(target string, graphClient *dgo.Dgraph) {
	for {
		q := `query allofterms($a: string) {
Node(func: allofterms(full, $a)) {
			 uid
			 nodeType
			 full
			 }

		 }`

		ctx := context.Background()
		txn := graphClient.NewTxn()
		res, err := txn.QueryWithVars(ctx, q, map[string]string{"$a": target})
		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:    target,
			Nodes:     []models.Node{},
			NodeType:  "matcher",
		}
		if len(n.Result) != 0 {
			time.Sleep(time.Duration(2) * time.Second)
			logrus.Info("Found paste match for ", target)
			// TODO: review time and id to be updated on new resulsts
			for _, res := range n.Result {
				if len(matcher.Nodes) == 0 {
					matcher.Nodes = append(matcher.Nodes, res)
					continue
				} else {
					for _, node := range matcher.Nodes {
						if res.UID != node.UID {
							matcher.Nodes = append(matcher.Nodes, res)

						}
					}
				}

			}

			query := fmt.Sprintf(`query { match as var(func: eq(target, "%s")) }`, target)

			pb, err := json.Marshal(models.Match{UID: "uid(match)", ID: matcher.ID, Target: target, Nodes: matcher.Nodes, NodeType: "matcher", Timestamp: rfc3339time})
			if err != nil {
				logrus.Fatal(err)
			}

			mu := &api.Mutation{
				SetJson: pb,
			}

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

			txn := graphClient.NewTxn()
			_, err = txn.Do(ctx, req)
			if err != nil {
				logrus.Warn(err)
			}
			// txn.Discard(ctx)
			time.Sleep(time.Duration(2) * time.Second)
		}

	}
}

func runCertstreamMatcher(target string, graphClient *dgo.Dgraph) {
	for {
		q := `query allofterms($a: string){
Node(func: allofterms(cn, $a)){
		fingerprint      
		notBefore        
		notAfter         
		cn               
		sourceName       
		basicConstraints 
		raw              
}
}`

		ctx := context.Background()
		txn := graphClient.NewTxn()
		res, err := txn.QueryWithVars(ctx, q, map[string]string{"$a": target})
		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:    target,
			Nodes:     []models.Node{},
			NodeType:  "matcher",
		}
		if len(n.Result) != 0 {
			time.Sleep(time.Duration(2) * time.Second)
			logrus.Info("Found certstream match for ", target)

			for _, res := range n.Result {
				if len(matcher.Nodes) == 0 {
					matcher.Nodes = append(matcher.Nodes, res)
					continue
				} else {
					for _, nodes := range matcher.Nodes {
						if res.UID != nodes.UID {
							matcher.Nodes = append(matcher.Nodes, res)
						}
					}
				}
			}

			query := fmt.Sprintf(`query { match as var(func: eq(target, "%s")) }`, target)

			pb, err := json.Marshal(models.Match{UID: "uid(match)", ID: matcher.ID, Target: target, Nodes: matcher.Nodes, NodeType: "matcher"})
			if err != nil {
				logrus.Fatal(err)
			}

			mu := &api.Mutation{
				SetJson: pb,
			}

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

			txn := graphClient.NewTxn()
			_, err = txn.Do(ctx, req)
			if err != nil {
				logrus.Warn(err)
			}

			time.Sleep(time.Duration(2) * time.Second)
		}
	}
}

func runShodanMatcher(target string, graphClient *dgo.Dgraph) {
	for {
		q := `query allofterms($a: string){
  Node(func: eq(hostnames, $a)){
    uid
    type
    hostnames
    product
    version
    title
    ip
    os
    organization
    isp
    cpe
    asn
    port
    html
    banner
    link
    transport
    domains
    timestamp
  }
}
`
		ctx := context.Background()
		txn := graphClient.NewTxn()
		res, err := txn.QueryWithVars(ctx, q, map[string]string{"$a": target})
		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:    target,
			Nodes:     []models.Node{},
			NodeType:  "matcher",
		}
		if len(n.Result) != 0 {
			time.Sleep(time.Duration(2) * time.Second)
			logrus.Info("Found shodan match for ", target)
			// TODO: review time and id to be updated on new resulsts

			for _, res := range n.Result {
				if len(matcher.Nodes) == 0 {
					matcher.Nodes = append(matcher.Nodes, res)
					continue
				} else {
					for _, node := range matcher.Nodes {
						if res.UID != node.UID {
							matcher.Nodes = append(matcher.Nodes, res)

						}
					}
				}

			}

			query := fmt.Sprintf(`query { match as var(func: eq(target, "%s")) }`, target)

			pb, err := json.Marshal(models.Match{UID: "uid(match)", ID: matcher.ID, Target: target, Nodes: matcher.Nodes, NodeType: "matcher", Timestamp: rfc3339time})
			if err != nil {
				logrus.Fatal(err)
			}

			mu := &api.Mutation{
				SetJson: pb,
			}

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

			txn := graphClient.NewTxn()
			_, err = txn.Do(ctx, req)
			if err != nil {
				logrus.Warn(err)
			}
			// txn.Discard(ctx)
			time.Sleep(time.Duration(2) * time.Second)

		}
	}
}

// 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
}