package main

import (
	"encoding/json"
	"github.com/go-chi/chi"
	"github.com/hillu/go-yara/v4"
	"github.com/package-url/packageurl-go"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"io"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

type callbackFunc func(yara.MatchRules, error)

type Job struct {
	Data     io.ReadCloser
	Callback callbackFunc
}

var (
	client *http.Client

	jobChan = make(chan *Job)

	jobCount      int64
	detectedCount int64
)

type statsResponse struct {
	Processed int64 `json:"jobsProcessed"`
	Detected  int64 `json:"detectionCount""`
}

type response struct {
	Success      bool            `json:"success"`
	Error        error           `json:"error,omitempty"`
	MatchedRules yara.MatchRules `json:"rules,omitempty"`
}

type multiResponse struct {
	Success bool                       `json:"success"`
	Files   map[string]yara.MatchRules `json:"files,omitempty"`
}

func main() {
	viper.SetDefault("threads", runtime.NumCPU())
	viper.SetDefault("rules", "pkg:github/Neo23x0/signature-base#yara")
	viper.SetDefault("bind", ":8080")

	viper.AutomaticEnv()

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

	c, err := yara.NewCompiler()

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

	c.DefineVariable("filename", "")
	c.DefineVariable("filepath", "")
	c.DefineVariable("extension", "")
	c.DefineVariable("filetype", "")

	log.Info("Loading rules")

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

	r := chi.NewRouter()

	r.Get("/stats", statsHandler)
	r.Post("/scan", scanHandler)

	bind := viper.GetString("bind")

	log.WithField("bind", bind).Info("Binding to address")

	go http.ListenAndServe(bind, r)

	ch := make(chan os.Signal)

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

	<-ch
}

func statsHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(statsResponse{Processed: atomic.LoadInt64(&jobCount), Detected: atomic.LoadInt64(&detectedCount)})
}

// HTTP handler for scanning files
func scanHandler(w http.ResponseWriter, r *http.Request) {
	contentType := r.Header.Get("Content-Type")

	if idx := strings.Index(contentType, ";"); idx != -1 {
		contentType = contentType[0:idx]
	}

	var res interface{}

	wg := &sync.WaitGroup{}

	switch contentType {
	case "multipart/form-data":
		if r.MultipartForm == nil {
			r.ParseMultipartForm(32 << 20)
		}

		log.WithField("contentType", contentType).Debug("Adding files from multipart form as jobs")

		fileCount := 0

		response := multiResponse{
			Success: true,
			Files:   make(map[string]yara.MatchRules),
		}

		fileLock := &sync.Mutex{}

		for _, files := range r.MultipartForm.File {
			// Append files
			for _, file := range files {
				wg.Add(1)

				fileCount++

				f, err := file.Open()

				if err != nil {
					continue
				}

				job := &Job{
					Data: f,
					Callback: func(m yara.MatchRules, err error) {
						fileLock.Lock()
						response.Files[file.Filename] = m
						fileLock.Unlock()

						wg.Done()
					},
				}

				jobChan <- job
			}
		}

		log.WithField("count", fileCount).Debug("Waiting for jobs to finish")

		wg.Wait()

		log.WithField("matches", res).Debug("Matched rules")

		res = response
	default:
		wg.Add(1)

		job := &Job{
			Data: r.Body,
			Callback: func(m yara.MatchRules, err error) {
				log.WithField("match", m).Debug("Matched rules")

				if err != nil {
					res = response{Success: false, Error: err}
				} else {
					res = response{Success: true, MatchedRules: m}
				}

				if err != nil {
					log.WithError(err).Error("Unable to send response")
				}

				wg.Done()
			},
		}

		log.WithField("contentType", contentType).Debug("Scanning contents of body")

		jobChan <- job
	}

	wg.Wait()

	w.Header().Set("Content-Type", "application/json")

	err := json.NewEncoder(w).Encode(res)

	if err != nil {
		log.WithError(err).Error("Unable to send response")
	}
}

// Load rules from the rules configuration option
func loadRules(c *yara.Compiler) {
	rulePaths := strings.Split(viper.GetString("rules"), ",")

	for _, p := range rulePaths {
		log.WithField("package", p).Info("Loading rules from package")

		instance, err := packageurl.FromString(p)

		if err != nil {
			log.WithFields(log.Fields{
				"error":   err,
				"package": p,
			}).Fatalln("Invalid rule URL")
		}

		switch instance.Type {
		case "git", "bitbucket", "github", "gitlab":
			err = loadRulesFromGit(instance, c)
		case "http", "https":
			err = loadRulesFromHttp(instance, c)
		}

		if err != nil {
			log.WithFields(log.Fields{
				"error":   err,
				"package": p,
			}).Fatalln("Unable to load rules")
		}
	}
}

// Load rules from http(s)
// TODO: Support archive files alongside standard yar files
func loadRulesFromHttp(pkg packageurl.PackageURL, c *yara.Compiler) error {
	res, err := client.Get(pkg.Type + "://" + pkg.Namespace + "/" + pkg.Name)

	if err != nil {
		return err
	}

	defer res.Body.Close()

	b, err := io.ReadAll(res.Body)

	if err != nil {
		return err
	}

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

// A worker routine. Creates a new scanner instance and pulls jobs.
func worker(rules *yara.Rules) {
	s, err := yara.NewScanner(rules)

	if err != nil {
		panic(err)
	}

	for {
		job := <-jobChan

		log.Debug("Processing job")

		processJob(s, job)
	}
}

func processJob(s *yara.Scanner, job *Job) {
	m := make(yara.MatchRules, 0)

	atomic.AddInt64(&jobCount, 1)

	defer job.Data.Close()

	b, err := io.ReadAll(job.Data)

	if err != nil {
		job.Callback(nil, err)
		return
	}

	err = s.SetCallback(&m).ScanMem(b)

	if err != nil {
		job.Callback(nil, err)
		return
	}

	if len(m) > 0 {
		atomic.AddInt64(&detectedCount, 1)
	}

	job.Callback(m, nil)
}