Disable debug, add stats, set content type on responses

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