package main

import (
	"archive/zip"
	"encoding/json"
	"errors"
	"github.com/beanstalkd/go-beanstalk"
	"github.com/hillu/go-yara/v4"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/afero"
	"github.com/spf13/afero/zipfs"
	"github.com/spf13/viper"
	"io"
	"io/fs"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path"
	"runtime"
	"strings"
	"syscall"
	"time"
)

type Job struct {
	PasteID string
	Data    []byte
}

var (
	client *http.Client
)

func main() {
	viper.SetDefault("beanstalk", "127.0.0.1:11300")
	viper.SetDefault("threads", runtime.NumCPU())
	viper.SetDefault("pasteUrl", "https://paste.ee")
	viper.SetDefault("rules", "https://github.com/Neo23x0/signature-base/yara")

	viper.AutomaticEnv()

	client = &http.Client{
		Timeout: 15 * time.Second,
	}

	jobChan := make(chan Job)

	c, err := yara.NewCompiler()

	if err != nil {
		log.WithError(err).Fatal("Unable to setup new compiler")
	}

	loadRules(c)

	rules, err := c.GetRules()

	if err != nil {
		log.WithError(err).Fatal("Unable to compile rules")
	}

	threads := viper.GetInt("threads")

	log.WithField("workers", threads).Info("Starting workers")

	for i := 0; i < threads; i++ {
		go worker(rules, jobChan)
	}

	conn, err := beanstalk.Dial("tcp", viper.GetString("beanstalk"))

	if err != nil {
		log.WithError(err).Fatal("Unable to connect to beanstalkd")
	}

	go watchQueue(conn, jobChan)

	ch := make(chan os.Signal)

	signal.Notify(ch, syscall.SIGKILL, syscall.SIGTERM, syscall.SIGINT)

	<-ch
}

func loadRules(c *yara.Compiler) {
	rulePaths := strings.Split(viper.GetString("rules"), ",")

	var err error
	var closer io.Closer

	for _, p := range rulePaths {
		var afs afero.Fs

		log.WithField("path", p).Info("Loading rules from p")

		if strings.HasPrefix(p, "http") || strings.HasPrefix(p, "git") {
			// load http rules
			afs, closer, err = loadRulesFromHttp(p)
		} else if stat, err := os.Stat(p); !os.IsNotExist(err) && stat.IsDir() {
			afs = afero.NewBasePathFs(afero.NewOsFs(), p)
		} else {
			continue
		}

		if err != nil {
			log.WithError(err).Fatalln("Unable to load rules")
		}

		afero.Walk(afs, "", func(path string, info fs.FileInfo, err error) error {
			b, err := afero.ReadFile(afs, path)

			return c.AddString(string(b), "")
		})

		if closer != nil {
			closer.Close()
		}
	}
}

func loadRulesFromHttp(ruleUrl string) (afero.Fs, io.Closer, error) {
	u, err := url.Parse(ruleUrl)

	if err != nil {
		log.WithError(err).Fatalln("Invalid URL!")
	}

	var subPath string

	if u.Host == "github.com" {
		parts := strings.Split(u.Path, "/")

		if len(parts) < 2 {
			return nil, nil, errors.New("invalid repo")
		}
		// User = parts[0], Repo = parts[1], Sub path = parts[2:]

		// https://github.com/parts[0]/parts[1]/archive/master.zip

		ruleUrl = "https://github.com/" + path.Join(parts[0], parts[1], "archive", "master.zip")
		subPath = path.Join(parts[2:]...)
	}

	res, err := client.Get(ruleUrl)

	if err != nil {
		return nil, nil, err
	}

	defer res.Body.Close()

	tmpFile, err := os.CreateTemp(os.TempDir(), "yararules.zip")

	io.Copy(tmpFile, res.Body)

	tmpFile.Seek(0, io.SeekStart)

	stat, err := tmpFile.Stat()

	if err != nil {
		return nil, nil, err
	}

	z, err := zip.NewReader(tmpFile, stat.Size())

	if err != nil {
		return nil, nil, err
	}

	return afero.NewBasePathFs(zipfs.New(z), subPath), tmpFile, nil
}

func watchQueue(c *beanstalk.Conn, jobChan chan Job) {
	for {
		id, body, err := c.Reserve(5 * time.Second)

		if err != nil {
			continue
		}

		log.WithFields(log.Fields{
			"id":   id,
			"body": body,
		}).Debug("Handling job")

		var job Job

		if err = json.Unmarshal(body, &job); err != nil {
			continue
		}

		jobChan <- job
	}
}

func worker(rules *yara.Rules, jobs chan Job) {
	s, err := yara.NewScanner(rules)

	if err != nil {
		panic(err)
	}

	for {
		job := <-jobs

		processJob(s, job)
	}
}

func processJob(s *yara.Scanner, job Job) {
	var m yara.MatchRules

	err := s.SetCallback(&m).ScanMem(job.Data)

	if err != nil {
		return
	}

	// Respond with job
	if len(m) < 1 {
		return
	}

	quarantine(job.PasteID, m[0].Rule)
}

func quarantine(pasteId, reason string) {
	v := make(url.Values)

	v.Set("reason", reason)

	req, err := http.NewRequest(http.MethodPost, viper.GetString("pasteUrl")+"/admin/quarantine/"+pasteId, nil)

	if err != nil {
		return
	}

	res, err := client.Do(req)

	if err != nil {
		return
	}

	defer res.Body.Close()
}