yarascanner/main.go

315 lines
5.9 KiB
Go

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