From 7fdac071d8c3c351b6381bc670961c47db8dc0fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pavel=20K=C3=A1cha?= <ph@cesnet.cz>
Date: Fri, 18 Feb 2022 17:21:19 +0100
Subject: [PATCH] Mitigate too big Idea events (close context if too big)

---
 contrib/connectors/hp-labrea/labrea-idea.py | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/contrib/connectors/hp-labrea/labrea-idea.py b/contrib/connectors/hp-labrea/labrea-idea.py
index a5316e1..f8b323a 100755
--- a/contrib/connectors/hp-labrea/labrea-idea.py
+++ b/contrib/connectors/hp-labrea/labrea-idea.py
@@ -34,6 +34,9 @@ class WindowContextMgr(object):
         self.ideagen = ideagen
         self.first_update_queue = OrderedDict()
         self.last_update_queue = OrderedDict()
+        # Hammer to mitigate too big events
+        self.max_count = 2000
+        self.max_src_ports = 1024
 
     def expire_queue(self, queue, window):
         aggr_events = []
@@ -68,9 +71,16 @@ class WindowContextMgr(object):
                     self.first_update_queue[ctx] = self.update_timestamp
                     self.last_update_queue[ctx] = self.update_timestamp
                 else:
-                    self.ctx_append(self.contexts[ctx], event)
-                    del self.last_update_queue[ctx]
-                    self.last_update_queue[ctx] = self.update_timestamp
+                    if not self.ctx_append(self.contexts[ctx], event):
+                        closed = self.ctx_close(self.contexts[ctx])
+                        if closed is not None:
+                            aggr_events.append(closed)
+                        del self.contexts[ctx]
+                        del self.first_update_queue[ctx]
+                        del self.last_update_queue[ctx]
+                    else:
+                        del self.last_update_queue[ctx]
+                        self.last_update_queue[ctx] = self.update_timestamp
 
         return aggr_events
 
@@ -107,6 +117,7 @@ class PingContextMgr(WindowContextMgr):
         ctx["tgt_ips"].add(event.tgt_ip)
         ctx["count"] += 1
         ctx["last_update"] = self.update_timestamp
+        return ctx["count"] < self.max_count
 
     def ctx_close(self, ctx):
         return self.ideagen.gen_idea(
@@ -143,11 +154,13 @@ class ConnectContextMgr(WindowContextMgr):
         ctx["src_ports"].add(event.src_port)
         ctx["count"] += 1
         ctx["last_update"] = self.update_timestamp
+        return ctx["count"] < self.max_count
 
     def ctx_close(self, ctx):
+        src_ports = ctx["src_ports"] if len(ctx["src_ports"]) <= self.max_src_ports else None
         return self.ideagen.gen_idea(
             src=ctx["src_ip"],
-            src_ports=ctx["src_ports"],
+            src_ports=src_ports,
             targets=ctx["tgt_ips_ports"].items(),
             detect_time=self.update_timestamp,
             event_time=ctx["first_update"],
-- 
GitLab