From 4c24c0f52803a6376239dfc267917997ca9e93b6 Mon Sep 17 00:00:00 2001
From: Tyler <tystuyfzand@gmail.com>
Date: Sun, 7 Mar 2021 21:26:57 -0500
Subject: [PATCH] Add success field, matched rules output

---
 main.go | 45 ++++++++++++++++++++++++++++-----------------
 1 file changed, 28 insertions(+), 17 deletions(-)

diff --git a/main.go b/main.go
index d4ad86f..92be0c2 100644
--- a/main.go
+++ b/main.go
@@ -31,6 +31,16 @@ var (
 	jobChan = make(chan *Job)
 )
 
+type response struct {
+	Success      bool            `json:"success"`
+	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")
@@ -106,22 +116,17 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 
 		wg := &sync.WaitGroup{}
 
-		results := make([]*yara.MatchRules, 0)
-
-		jobCallback := func(m *yara.MatchRules, err error) {
-			wg.Done()
-
-			if err != nil {
-				return
-			}
-
-			results = append(results, m)
-		}
-
 		log.WithField("contentType", contentType).Debug("Adding files from multipart form as jobs")
 
 		fileCount := 0
 
+		res := multiResponse{
+			Success: true,
+			Files:   make(map[string]yara.MatchRules),
+		}
+
+		fileLock := &sync.Mutex{}
+
 		for _, files := range r.MultipartForm.File {
 			// Append files
 			for _, file := range files {
@@ -136,8 +141,14 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 				}
 
 				job := &Job{
-					Data:     f,
-					Callback: jobCallback,
+					Data: f,
+					Callback: func(m *yara.MatchRules, err error) {
+						fileLock.Lock()
+						res.Files[file.Filename] = *m
+						fileLock.Unlock()
+
+						wg.Done()
+					},
 				}
 
 				jobChan <- job
@@ -148,8 +159,8 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 
 		wg.Wait()
 
-		log.WithField("matches", w).Debug("Matched rules")
-		json.NewEncoder(w).Encode(results)
+		log.WithField("matches", res).Debug("Matched rules")
+		json.NewEncoder(w).Encode(res)
 	default:
 		job := &Job{
 			Data: r.Body,
@@ -160,7 +171,7 @@ func scanHandler(w http.ResponseWriter, r *http.Request) {
 					w.WriteHeader(http.StatusInternalServerError)
 					w.Write([]byte(err.Error()))
 				} else {
-					json.NewEncoder(w).Encode(m)
+					json.NewEncoder(w).Encode(response{Success: true, MatchedRules: *m})
 				}
 			},
 		}