Add success field, matched rules output
All checks were successful
continuous-integration/drone/push Build is passing

This commit is contained in:
Tyler 2021-03-07 21:26:57 -05:00
parent 6218a8ca50
commit 4c24c0f528

45
main.go
View File

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