From ba7a5a3afe73c00066c8473b453b0279df84249c Mon Sep 17 00:00:00 2001
From: Emanuel Seemann <3380606+seemanne@users.noreply.github.com>
Date: Wed, 20 Sep 2023 14:28:11 +0200
Subject: [PATCH] expose metrics to file

---
 pkg/bayesiantrain/inference.go | 40 ++++++++++++++++++++++++++++++++--
 pkg/bayesiantrain/trainer.go   |  8 ++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/pkg/bayesiantrain/inference.go b/pkg/bayesiantrain/inference.go
index 9ae349e00..5cdc5c179 100644
--- a/pkg/bayesiantrain/inference.go
+++ b/pkg/bayesiantrain/inference.go
@@ -2,6 +2,7 @@ package bayesiantrain
 
 import (
 	"fmt"
+	"os"
 
 	"github.com/antonmedv/expr"
 	"github.com/antonmedv/expr/vm"
@@ -15,7 +16,14 @@ type fakeBucket struct {
 	label  int
 }
 
-func (f *fakeBucket) scoreTrainedClassifier(results map[string]BayesianResult, exprCache map[string]vm.Program, prior float32, threshold float32) int {
+type inferenceResult struct {
+	ip          string
+	prediction  int
+	label       int
+	probability float32
+}
+
+func (f *fakeBucket) scoreTrainedClassifier(results map[string]BayesianResult, exprCache map[string]vm.Program, prior float32, threshold float32, resultChan chan<- inferenceResult) int {
 	var posterior float32
 	var queue leakybucket.Queue
 	var program vm.Program
@@ -23,6 +31,8 @@ func (f *fakeBucket) scoreTrainedClassifier(results map[string]BayesianResult, e
 	var ok bool
 	var guillotinecache map[string]bool
 
+	ip := f.events[0].Meta["source_ip"]
+	label := f.label
 	guillotinecache = make(map[string]bool)
 
 	for index, evt := range f.events {
@@ -60,9 +70,35 @@ func (f *fakeBucket) scoreTrainedClassifier(results map[string]BayesianResult, e
 		}
 
 		if posterior >= threshold {
+			resultChan <- inferenceResult{ip, 1, label, posterior}
 			return 1
 		}
 	}
-
+	resultChan <- inferenceResult{ip, 0, label, posterior}
 	return 0
 }
+
+func saveResultsToDisk(inputChan <-chan inferenceResult) {
+	var res inferenceResult
+	var str string
+	var more bool
+
+	f, err := os.Create("inference_result.csv")
+
+	if err != nil {
+		fmt.Printf("%s", err)
+	}
+
+	f.WriteString("ip,probability,label\n")
+
+	defer f.Close()
+
+	for {
+		res, more = <-inputChan
+		if !more {
+			return
+		}
+		str = fmt.Sprint(res.ip, ",", res.probability, ",", res.label, "\n")
+		f.WriteString(str)
+	}
+}
diff --git a/pkg/bayesiantrain/trainer.go b/pkg/bayesiantrain/trainer.go
index b5466a7d5..f783b06a7 100644
--- a/pkg/bayesiantrain/trainer.go
+++ b/pkg/bayesiantrain/trainer.go
@@ -143,12 +143,16 @@ func (s *LogEventStorage) GenerateBucketMetrics(threshold float32) error {
 	var falseNegative int
 	var truePositive int
 	var trueNegative int
+	var inferenceChannel chan inferenceResult
 
+	inferenceChannel = make(chan inferenceResult, 20)
 	prior := float32(s.nEvilIps) / float32(s.total)
 
+	go saveResultsToDisk(inferenceChannel)
+
 	for _, bucket := range s.ParsedIpEvents {
 
-		res = bucket.scoreTrainedClassifier(s.CachedHypotheses, s.exprCache, prior, threshold)
+		res = bucket.scoreTrainedClassifier(s.CachedHypotheses, s.exprCache, prior, threshold, inferenceChannel)
 		if res < 0 {
 			return fmt.Errorf("generatebucketmetrics returned an error, aborting")
 		}
@@ -167,6 +171,8 @@ func (s *LogEventStorage) GenerateBucketMetrics(threshold float32) error {
 		}
 	}
 
+	close(inferenceChannel)
+
 	fmt.Println("raw : ", falsePositive, falseNegative, truePositive, trueNegative)
 	printBucketMetrics(falsePositive, falseNegative, truePositive, trueNegative)