Add caching mechanism and re-enable threat intelligence queue

Implemented a cache for blocklist checks to improve performance and reduced redundant processing. Re-enabled the threat intelligence queueing mechanism to ensure logs are properly analyzed and flagged. Additionally, cleaned up unused code and improved the blocklisted element handling logic.

Author: osmontero

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,7 +75,7 @@ func NewLog(c *gin.Context) {
 	}
 
 	cache.AddToCache(l)
-	//ti.Enqueue(l)
+	ti.Enqueue(l)
 	search.AddToQueue(l)
 	response["status"] = "queued"
 	c.JSON(http.StatusOK, response)

correlation/ti/bases.go

@@ -12,8 +12,6 @@ func Load() {
 
 	var files = []string{
 		"ip_blocklist.list",
-		//"domain_blocklist.list",
-		//"hostname_blocklist.list",
 	}
 
 	for _, file := range files {
@@ -22,10 +20,6 @@ func Load() {
 		switch file {
 		case "ip_blocklist.list":
 			t = "IP"
-			//case "domain_blocklist.list":
-			//	t = "domain"
-			//case "hostname_blocklist.list":
-			//	t = "hostname"
 		}
 
 		f, err := os.Open(filepath.Join("/app", file))

correlation/ti/ti.go

@@ -5,14 +5,68 @@ import (
 	"github.com/utmstack/UTMStack/correlation/correlation"
 	"github.com/utmstack/UTMStack/correlation/utils"
 	"runtime"
+	"strings"
+	"sync"
+	"time"
 )
 
+type Cache map[string]string
+
 var blockList map[string]string
 var channel chan string
+var cache Cache
+var cacheLock *sync.RWMutex
 
 func init() {
 	blockList = make(map[string]string, 10000)
 	channel = make(chan string, 10000)
+	cache = make(Cache, 10000)
+	cacheLock = &sync.RWMutex{}
+	cache.AutoClean()
+}
+
+func blocked(log string) bool {
+	log = strings.ToLower(log)
+
+	exclusionList := []string{
+		"block",
+		"denied",
+		"drop",
+		"reject",
+		"deny",
+	}
+
+	for _, e := range exclusionList {
+		if strings.Contains(log, e) {
+			return true
+		}
+	}
+
+	return false
+}
+
+func (c *Cache) Add(k, v string) {
+	cacheLock.Lock()
+	defer cacheLock.Unlock()
+	(*c)[k] = v
+}
+
+func (c *Cache) IsCached(k string) bool {
+	cacheLock.RLock()
+	defer cacheLock.RUnlock()
+	_, ok := (*c)[k]
+	return ok
+}
+
+func (c *Cache) AutoClean() {
+	go func() {
+		for {
+			cacheLock.Lock()
+			cache = make(Cache, 10000)
+			cacheLock.Unlock()
+			time.Sleep(8 * time.Hour)
+		}
+	}()
 }
 
 func IsBlocklisted() {
@@ -48,53 +102,30 @@ func IsBlocklisted() {
 				sourceIp := gjson.Get(log, "logx.*.src_ip")
 				destinationIp := gjson.Get(log, "logx.*.dest_ip")
 
-				for key, value := range blockList {
-					var stop bool
-
-					switch value {
-					case "IP":
-						if sourceIp.String() == key {
-							correlation.Alert(
-								"Connection attempt from a malicious IP",
-								"Low",
-								"A blocklisted element has been identified in the logs. Further investigation is recommended.",
-								"",
-								"Threat Intelligence",
-								"",
-								[]string{"https://threatwinds.com"},
-								gjson.Get(log, "dataType").String(),
-								gjson.Get(log, "dataSource").String(),
-								utils.ExtractDetails(saveFields, log),
-							)
-
-							stop = true
-						}
-
-						if destinationIp.String() == key {
-							correlation.Alert(
-								"Connection attempt to a malicious IP",
-								"High",
-								"A blocklisted element has been identified in the logs. Further investigation is recommended.",
-								"",
-								"Threat Intelligence",
-								"",
-								[]string{"https://threatwinds.com"},
-								gjson.Get(log, "dataType").String(),
-								gjson.Get(log, "dataSource").String(),
-								utils.ExtractDetails(saveFields, log),
-							)
-
-							stop = true
-						}
-					}
+				if !cache.IsCached(sourceIp.String()) {
+					if _, ok := blockList[sourceIp.String()]; ok && !blocked(log) {
+						correlation.Alert(
+							"Connection attempt from a malicious IP",
+							"Low",
+							"A blocklisted element has been identified in the logs. Further investigation is recommended.",
+							"",
+							"Threat Intelligence",
+							"",
+							[]string{"https://threatwinds.com"},
+							gjson.Get(log, "dataType").String(),
+							gjson.Get(log, "dataSource").String(),
+							utils.ExtractDetails(saveFields, log),
+						)
 
-					if stop {
-						break
 					}
 
-					/*if strings.Contains(log, key) {
+					cache.Add(sourceIp.String(), "true")
+				}
+
+				if !cache.IsCached(destinationIp.String()) {
+					if _, ok := blockList[destinationIp.String()]; ok && !blocked(log) {
 						correlation.Alert(
-							fmt.Sprintf("Malicious %s found in log: %s", value, key),
+							"Connection attempt from a malicious IP",
 							"Low",
 							"A blocklisted element has been identified in the logs. Further investigation is recommended.",
 							"",
@@ -105,15 +136,19 @@ func IsBlocklisted() {
 							gjson.Get(log, "dataSource").String(),
 							utils.ExtractDetails(saveFields, log),
 						)
+					}
 
-						break
-					}*/
+					cache.Add(destinationIp.String(), "true")
 				}
 			}
 		}()
 	}
 }
 
 func Enqueue(log string) {
+	if len(channel) >= 10000 {
+		return
+	}
+
 	channel <- log
 }