diff --git a/categorizer/analysis/ChromaAnalyser.go b/categorizer/analysis/ChromaAnalyser.go index a80335b..b826290 100644 --- a/categorizer/analysis/ChromaAnalyser.go +++ b/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 diff --git a/categorizer/analysis/StaticAnalyser.py b/categorizer/analysis/StaticAnalyser.py new file mode 100755 index 0000000..b31d4ae --- /dev/null +++ b/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() \ No newline at end of file diff --git a/categorizer/categorizer.go b/categorizer/categorizer.go index 62a3729..2b0692d 100644 --- a/categorizer/categorizer.go +++ b/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) diff --git a/categorizer/config/ConfigHandler.go b/categorizer/config/ConfigHandler.go index dd8cf87..1982414 100644 --- a/categorizer/config/ConfigHandler.go +++ b/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 { diff --git a/categorizer/config/config.json b/categorizer/config/config.json index a93a1a7..f0df359 100644 --- a/categorizer/config/config.json +++ b/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 - } -} \ No newline at end of file +} diff --git a/categorizer/controllers/AnalysisController.go b/categorizer/controllers/AnalysisController.go index 23ba1c5..5bbd334 100644 --- a/categorizer/controllers/AnalysisController.go +++ b/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) + } } } } diff --git a/categorizer/tests/CaronteRetriever_test.go b/categorizer/tests/CaronteRetriever_test.go index 881853e..b420295 100644 --- a/categorizer/tests/CaronteRetriever_test.go +++ b/categorizer/tests/CaronteRetriever_test.go @@ -10,12 +10,12 @@ 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 { @@ -23,7 +23,7 @@ func TestCaronteRetriever(t *testing.T) { t.Log(result) } } - ctx.Done() + cancel() close(queue) return } diff --git a/categorizer/tests/ConfigHandler_test.go b/categorizer/tests/ConfigHandler_test.go new file mode 100644 index 0000000..cce621a --- /dev/null +++ b/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") + } +} diff --git a/categorizer/tests/StaticAnalyser_test.go b/categorizer/tests/StaticAnalyser_test.go index 3fc6446..281c285 100644 --- a/categorizer/tests/StaticAnalyser_test.go +++ b/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: "", 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 } diff --git a/categorizer/tests/TulipRetriever_test.go b/categorizer/tests/TulipRetriever_test.go index cc67483..f041671 100644 --- a/categorizer/tests/TulipRetriever_test.go +++ b/categorizer/tests/TulipRetriever_test.go @@ -10,12 +10,12 @@ 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 { @@ -23,7 +23,7 @@ func TestTulipRetriever(t *testing.T) { t.Log(result) } } - ctx.Done() + cancel() close(queue) return }