first commit

main.go

@@ -0,0 +1,244 @@
+package main
+
+import (
+	"archive/zip"
+	"encoding/json"
+	"errors"
+	"github.com/beanstalkd/go-beanstalk"
+	"github.com/hillu/go-yara"
+	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) {
+	matched, err := s.ScanMem(job.Data)
+
+	if err != nil {
+		return
+	}
+
+	// Respond with job
+	if len(matched) < 1 {
+		return
+	}
+
+	quarantine(job.PasteID, matched[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()
+}