diff --git a/main.go b/main.go index 19d55ca..d97324f 100644 --- a/main.go +++ b/main.go @@ -14,6 +14,7 @@ import ( "runtime" "strings" "sync" + "sync/atomic" "syscall" "time" ) @@ -29,8 +30,16 @@ 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"` @@ -53,8 +62,6 @@ func main() { Timeout: 15 * time.Second, } - log.SetLevel(log.DebugLevel) - c, err := yara.NewCompiler() if err != nil { @@ -86,6 +93,7 @@ func main() { r := chi.NewRouter() + r.Get("/stats", statsHandler) r.Post("/scan", scanHandler) bind := viper.GetString("bind") @@ -101,6 +109,11 @@ func main() { <-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") @@ -194,6 +207,8 @@ func scanHandler(w http.ResponseWriter, r *http.Request) { wg.Wait() + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(res) if err != nil { @@ -273,6 +288,8 @@ func worker(rules *yara.Rules) { 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) @@ -289,5 +306,9 @@ func processJob(s *yara.Scanner, job *Job) { return } + if len(m) > 0 { + atomic.AddInt64(&detectedCount, 1) + } + job.Callback(m, nil) }