Refactor and optimize geo and TI processing.

Reorganized GeoIP database and threat intelligence loading into more modular functions for improved maintainability and code readability. Simplified caching, removed unused database function, and restructured rule-handling logic. Addressed minor variable renames and logging adjustments for consistency.

Author: osmontero

correlation/Dockerfile

@@ -1,6 +1,6 @@
-FROM ubuntu:22.04
-RUN apt update
-RUN apt install -y ca-certificates git
+FROM ubuntu:24.04
+RUN apt-get update
+RUN apt-get install -y ca-certificates git wget
 COPY correlation /app/
 COPY docs/swagger.json /app/docs/
 COPY docs/swagger.yaml /app/docs/
@@ -9,4 +9,13 @@ COPY run.sh /
 RUN chmod +x /app/correlation
 RUN chmod +x /run.sh
 RUN update-ca-certificates
+RUN wget -O /app/asn-blocks-v4.csv https://cdn.utmstack.com/geoip/asn-blocks-v4.csv
+RUN wget -O /app/asn-blocks-v6.csv https://cdn.utmstack.com/geoip/asn-blocks-v6.csv
+RUN wget -O /app/blocks-v4.csv https://cdn.utmstack.com/geoip/blocks-v4.csv
+RUN wget -O /app/blocks-v6.csv https://cdn.utmstack.com/geoip/blocks-v6.csv
+RUN wget -O /app/locations-en.csv https://cdn.utmstack.com/geoip/locations-en.csv
+RUN wget -O /app/ip_blocklist.list https://intelligence.threatwinds.com/feeds/public/ip/cumulative.list
+RUN wget -O /app/domain_blocklist.list https://intelligence.threatwinds.com/feeds/public/domain/cumulative.list
+RUN wget -O /app/hostname_blocklist.list https://intelligence.threatwinds.com/feeds/public/hostname/cumulative.list
+
 ENTRYPOINT [ "/run.sh" ]

correlation/api/newLogHandler.go

@@ -3,6 +3,7 @@ package api
 import (
 	"encoding/json"
 	"fmt"
+	"github.com/utmstack/UTMStack/correlation/ti"
 	"io"
 	"log"
 	"net/http"
@@ -74,6 +75,7 @@ func NewLog(c *gin.Context) {
 	}
 
 	cache.AddToCache(l)
+	ti.Enqueue(l)
 	search.AddToQueue(l)
 	response["status"] = "queued"
 	c.JSON(http.StatusOK, response)

correlation/cache/cache.go

@@ -13,15 +13,15 @@ import (
 
 const bufferSize int = 1000000
 
-var cacheStorageMutex = &sync.RWMutex{}
+var storageMutex = &sync.RWMutex{}
 
-var CacheStorage []string
+var storage []string
 
 func Status() {
 	for {
-		log.Printf("Logs in cache: %v", len(CacheStorage))
-		if len(CacheStorage) != 0 {
-			est := gjson.Get(CacheStorage[0], "@timestamp").String()
+		log.Printf("Logs in cache: %v", len(storage))
+		if len(storage) != 0 {
+			est := gjson.Get(storage[0], "@timestamp").String()
 			log.Printf("Old document in cache: %s", est)
 		}
 		time.Sleep(60 * time.Second)
@@ -31,8 +31,8 @@ func Status() {
 
 func Search(allOf []rules.AllOf, oneOf []rules.OneOf, seconds int64) []string {
 	var elements []string
-	cacheStorageMutex.RLock()
-	defer cacheStorageMutex.RUnlock()
+	storageMutex.RLock()
+	defer storageMutex.RUnlock()
 
 	cToBreak := 0
 	ait := time.Now().UTC().Unix() - func() int64 {
@@ -43,8 +43,8 @@ func Search(allOf []rules.AllOf, oneOf []rules.OneOf, seconds int64) []string {
 			return seconds
 		}
 	}()
-	for i := len(CacheStorage) - 1; i >= 0; i-- {
-		est := gjson.Get(CacheStorage[i], "@timestamp").String()
+	for i := len(storage) - 1; i >= 0; i-- {
+		est := gjson.Get(storage[i], "@timestamp").String()
 		eit, err := time.Parse(time.RFC3339Nano, est)
 		if err != nil {
 			log.Printf("Could not parse @timestamp: %v", err)
@@ -61,23 +61,23 @@ func Search(allOf []rules.AllOf, oneOf []rules.OneOf, seconds int64) []string {
 			var allCatch bool
 			var oneCatch bool
 			for _, of := range oneOf {
-				oneCatch = evalElement(CacheStorage[i], of.Field, of.Operator, of.Value)
+				oneCatch = evalElement(storage[i], of.Field, of.Operator, of.Value)
 				if oneCatch {
 					break
 				}
 			}
 			for _, af := range allOf {
-				allCatch = evalElement(CacheStorage[i], af.Field, af.Operator, af.Value)
+				allCatch = evalElement(storage[i], af.Field, af.Operator, af.Value)
 				if !allCatch {
 					break
 				}
 			}
 			if (len(allOf) == 0 || allCatch) && (len(oneOf) == 0 || oneCatch) {
-				elements = append(elements, CacheStorage[i])
+				elements = append(elements, storage[i])
 			}
 		}
 	}
-	
+
 	return elements
 }
 
@@ -97,9 +97,9 @@ func ProcessQueue() {
 		go func() {
 			for {
 				l := <-logs
-				cacheStorageMutex.Lock()
-				CacheStorage = append(CacheStorage, l)
-				cacheStorageMutex.Unlock()
+				storageMutex.Lock()
+				storage = append(storage, l)
+				storageMutex.Unlock()
 			}
 		}()
 	}
@@ -109,11 +109,11 @@ func Clean() {
 	for {
 		var clean bool
 
-		if len(CacheStorage) > 1 {
+		if len(storage) > 1 {
 			if utils.AssignedMemory >= 80 {
 				clean = true
 			} else {
-				old := gjson.Get(CacheStorage[0], "@timestamp").String()
+				old := gjson.Get(storage[0], "@timestamp").String()
 				oldTime, err := time.Parse(time.RFC3339Nano, old)
 				if err != nil {
 					log.Printf("Could not parse old log timestamp. Cleaning up")
@@ -129,9 +129,9 @@ func Clean() {
 		}
 
 		if clean {
-			cacheStorageMutex.Lock()
-			CacheStorage = CacheStorage[1:]
-			cacheStorageMutex.Unlock()
+			storageMutex.Lock()
+			storage = storage[1:]
+			storageMutex.Unlock()
 		} else {
 			time.Sleep(5 * time.Second)
 		}

correlation/cache/cache_test.go

@@ -1,11 +1,8 @@
-package cache_test
+package cache
 
 import (
-	"testing"
-	
-
-	"github.com/utmstack/UTMStack/correlation/cache"
 	"github.com/utmstack/UTMStack/correlation/rules"
+	"testing"
 )
 
 func TestSearch(t *testing.T) {
@@ -16,7 +13,7 @@ func TestSearch(t *testing.T) {
 		`{"@timestamp":"2022-01-01T00:00:03.000Z","field1":"value1","field2":"value2"}`,
 		`{"@timestamp":"2022-01-01T00:00:04.000Z","field1":"value1","field2":"value2"}`,
 	}
-	cache.CacheStorage = cacheStorage
+	storage = cacheStorage
 	allOf := []rules.AllOf{
 		{Field: "field1", Operator: "==", Value: "value1"},
 	}
@@ -31,7 +28,7 @@ func TestSearch(t *testing.T) {
 		`{"@timestamp":"2022-01-01T00:00:01.000Z","field1":"value1","field2":"value2"}`,
 		`{"@timestamp":"2022-01-01T00:00:00.000Z","field1":"value1","field2":"value2"}`,
 	}
-	result := cache.Search(allOf, oneOf, int64(seconds))
+	result := Search(allOf, oneOf, int64(seconds))
 	if len(result) != len(expected) {
 		t.Errorf("Expected %d elements, but got %d", len(expected), len(result))
 	}
@@ -40,4 +37,4 @@ func TestSearch(t *testing.T) {
 			t.Errorf("Expected %s, but got %s", expected[i], r)
 		}
 	}
-}
+}

correlation/correlation/analyzer.go

@@ -1,85 +1,25 @@
 package correlation
 
 import (
-	"net"
-
-	"github.com/utmstack/UTMStack/correlation/geo"
-	"github.com/utmstack/UTMStack/correlation/rules"
 	"github.com/tidwall/gjson"
+	"github.com/utmstack/UTMStack/correlation/rules"
+	"github.com/utmstack/UTMStack/correlation/utils"
 )
 
-func processResponse(logs []string, rule rules.Rule, save []rules.SavedField, tmpLogs *[20][]map[string]string,
+func processResponse(logs []string, rule rules.Rule, save []utils.SavedField, tmpLogs *[20][]map[string]string,
 	steps, step, minCount int) {
-	if len(logs) >= func()int{
-		switch minCount{
+	if len(logs) >= func() int {
+		switch minCount {
 		case 0:
 			return 1
 		default:
 			return minCount
 		}
 	}() {
 		for _, l := range logs {
-			var fields = map[string]string{
-				"id": gjson.Get(l, "id").String(),
-			}
-			//User saved fields
-			for _, save := range save {
-				fields[save.Alias] = gjson.Get(l, save.Field).String()
-			}
-			// Try to resolve SourceHost if SourceIP exists but not SourceHost
-			if fields["SourceHost"] == "" && fields["SourceIP"] != "" {
-				host, _ := net.LookupHost(fields["SourceIP"])
-				fields["SourceHost"] = host[0]
-			}
-			// Try to resolve DestinationHost if DestinationIP exists but not DestinationHost
-			if fields["DestinationHost"] == "" && fields["DestinationIP"] != "" {
-				host, _ := net.LookupHost(fields["DestinationIP"])
-				fields["DestinationHost"] = host[0]
-			}
-			// Try to resolve SourceIP if SourceHost exists but not SourceIP
-			if fields["SourceHost"] != "" && fields["SourceIP"] == "" {
-				ip, _ := net.LookupIP(fields["SourceHost"])
-				if len(ip) != 0 && ip[0].String() != "<nil>" {
-					fields["SourceIP"] = ip[0].String()
-				}
-			}
-			// Try to resolve DestinationIP if DestinationHost exists but not DestinationIP
-			if fields["DestinationHost"] != "" && fields["DestinationIP"] == "" {
-				ip, _ := net.LookupIP(fields["DestinationHost"])
-				if len(ip) != 0 && ip[0].String() != "<nil>" {
-					fields["DestinationIP"] = ip[0].String()
-				}
-			}
-			// Try to geolocate SourceIP if exists
-			if fields["SourceIP"] != "" {
-				location := geo.Geolocate(fields["SourceIP"])
-				fields["SourceCountry"] = location["country"]
-				fields["SourceCountryCode"] = location["countryCode"]
-				fields["SourceCity"] = location["city"]
-				fields["SourceLat"] = location["latitude"]
-				fields["SourceLon"] = location["longitude"]
-				fields["SourceAccuracyRadius"] = location["accuracyRadius"]
-				fields["SourceASN"] = location["asn"]
-				fields["SourceASO"] = location["aso"]
-				fields["SourceIsSatelliteProvider"] = location["isSatelliteProvider"]
-				fields["SourceIsAnonymousProxy"] = location["isAnonymousProxy"]
-			}
-			// Try to geolocate DetinationIP if exists
-			if fields["DestinationIP"] != "" {
-				location := geo.Geolocate(fields["DestinationIP"])
-				fields["DestinationCountry"] = location["country"]
-				fields["DestinationCountryCode"] = location["countryCode"]
-				fields["DestinationCity"] = location["city"]
-				fields["DestinationLat"] = location["latitude"]
-				fields["DestinationLon"] = location["longitude"]
-				fields["DestinationAccuracyRadius"] = location["accuracyRadius"]
-				fields["DestinationASN"] = location["asn"]
-				fields["DestinationASO"] = location["aso"]
-				fields["DestinationIsSatelliteProvider"] = location["isSatelliteProvider"]
-				fields["DestinationIsAnonymousProxy"] = location["isAnonymousProxy"]
-			}
+			fields := utils.ExtractDetails(save, l)
 
-			// Alert in the last step or save data to next cicle
+			// Alert in the last step or save data to the next iteration
 			if steps-1 == step {
 				// Use content of AlertName as Name if exists
 				var alertName string

correlation/correlation/reporter.go

@@ -2,16 +2,15 @@ package correlation
 
 import (
 	"encoding/json"
-	"strconv"
-	"strings"
-	"time"
-
-	"log"
-
 	"github.com/google/uuid"
 	"github.com/levigross/grequests"
+	"github.com/utmstack/UTMStack/correlation/geo"
 	"github.com/utmstack/UTMStack/correlation/search"
 	"github.com/utmstack/UTMStack/correlation/utils"
+	"log"
+	"strconv"
+	"strings"
+	"time"
 )
 
 type Host struct {
@@ -65,13 +64,47 @@ type AlertFields struct {
 }
 
 func Alert(name, severity, description, solution, category, tactic string, reference []string, dataType, dataSource string,
-	details map[string]string) {
+	fields map[string]string) {
+
+	// Try to geolocate SourceIP if exists
+	if fields["SourceIP"] != "" {
+		location := geo.Geolocate(fields["SourceIP"])
+		if len(location) != 0 {
+			fields["SourceCountry"] = location["country"]
+			fields["SourceCountryCode"] = location["countryCode"]
+			fields["SourceCity"] = location["city"]
+			fields["SourceLat"] = location["latitude"]
+			fields["SourceLon"] = location["longitude"]
+			fields["SourceAccuracyRadius"] = location["accuracyRadius"]
+			fields["SourceASN"] = location["asn"]
+			fields["SourceASO"] = location["aso"]
+			fields["SourceIsSatelliteProvider"] = location["isSatelliteProvider"]
+			fields["SourceIsAnonymousProxy"] = location["isAnonymousProxy"]
+		}
+	}
+
+	// Try to geolocate DestinationIP if exists
+	if fields["DestinationIP"] != "" {
+		location := geo.Geolocate(fields["DestinationIP"])
+		if len(location) != 0 {
+			fields["DestinationCountry"] = location["country"]
+			fields["DestinationCountryCode"] = location["countryCode"]
+			fields["DestinationCity"] = location["city"]
+			fields["DestinationLat"] = location["latitude"]
+			fields["DestinationLon"] = location["longitude"]
+			fields["DestinationAccuracyRadius"] = location["accuracyRadius"]
+			fields["DestinationASN"] = location["asn"]
+			fields["DestinationASO"] = location["aso"]
+			fields["DestinationIsSatelliteProvider"] = location["isSatelliteProvider"]
+			fields["DestinationIsAnonymousProxy"] = location["isAnonymousProxy"]
+		}
+	}
 
 	log.Printf("Reporting alert: %s", name)
 
-	if !UpdateAlert(name, severity, details) {
+	if !UpdateAlert(name, severity, fields) {
 		NewAlert(name, severity, description, solution, category, tactic, reference, dataType, dataSource,
-			details)
+			fields)
 	}
 }
 
@@ -208,11 +241,12 @@ func UpdateAlert(name, severity string, details map[string]string) bool {
 							},
 						},
 					})
-					_ = r.Close()
 					if err != nil {
 						log.Printf("Could not update existent alert: %v", err)
 						return false
 					}
+
+					_ = r.Close()
 				}
 			}
 		}

correlation/docs/docs.go

@@ -10,11 +10,11 @@ const docTemplate = `{
         "description": "{{escape .Description}}",
         "title": "{{.Title}}",
         "contact": {
-            "name": "Osmany Montero",
-            "email": "osmany@quantfall.com"
+            "name": "UTMStack LLC",
+            "email": "contact@utmstack.com"
         },
         "license": {
-            "name": "Private"
+            "name": "AGPLv3"
         },
         "version": "{{.Version}}"
     },
@@ -47,7 +47,7 @@ var SwaggerInfo = &swag.Spec{
 	BasePath:         "/v1",
 	Schemes:          []string{},
 	Title:            "UTMStack's Correlation Engine",
-	Description:      "Rules based correlation engine for UTMStack.",
+	Description:      "Rules-based correlation engine for UTMStack.",
 	InfoInstanceName: "swagger",
 	SwaggerTemplate:  docTemplate,
 	LeftDelim:        "{{",