From 2356c6bcd7a9760b78b1ed3766c3fdbdab920c93 Mon Sep 17 00:00:00 2001
From: Mark Veidemanis <m@zm.is>
Date: Thu, 9 Feb 2023 22:54:38 +0000
Subject: [PATCH] Fix formatting matched arguments

---
 core/db/elastic.py |  2 ++
 core/lib/rules.py  | 30 ++++++++++++++++++++++++++----
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/core/db/elastic.py b/core/db/elastic.py
index b034a5c..3d4151a 100644
--- a/core/db/elastic.py
+++ b/core/db/elastic.py
@@ -371,6 +371,8 @@ class ElasticsearchBackend(StorageBackend):
         """
         Check the results of a scheduled query for aggregations.
         """
+        if rule_object.aggs is None:
+            return result_map
         for index, (meta, result) in result_map.items():
             # Default to true, if no aggs are found, we still want to match
             match = True
diff --git a/core/lib/rules.py b/core/lib/rules.py
index 0048cc8..e00d02d 100644
--- a/core/lib/rules.py
+++ b/core/lib/rules.py
@@ -103,7 +103,7 @@ def format_webhook(**kwargs):
     notify_message = {
         "rule_id": rule.id,
         "rule_name": rule.name,
-        "match": matched,
+        "matched": matched,
         "total_hits": total_hits,
         "index": index,
         "data": message,
@@ -215,6 +215,25 @@ class NotificationRuleData(object):
                 self.object.match[index] = False
         self.object.save()
 
+    def format_matched(self, messages):
+        matched = {}
+        for message in messages:
+            for field, value in self.data:
+                if field == "msg":
+                    # Allow partial matches for msg
+                    for msg in value:
+                        if "msg" in message:
+                            if msg.lower() in message["msg"].lower():
+                                matched[field] = msg
+                                # Break out of the msg matching loop
+                                break
+                    # Continue to next field
+                    continue
+                if field in message and message[field] in value:
+                    # Do exact matches for all other fields
+                    matched[field] = message[field]
+        return matched
+
     def store_match(self, index, match):
         """
         Store a match result.
@@ -293,8 +312,9 @@ class NotificationRuleData(object):
         for agg_name, agg in aggs.items():
             print("ITER", agg_name, agg)
             # Already checked membership below
-            op, value = self.aggs[agg_name]
-            new_aggs[agg_name] = f"{agg['value']}{op}{value}"
+            if agg_name in self.aggs:
+                op, value = self.aggs[agg_name]
+                new_aggs[agg_name] = f"{agg['value']}{op}{value}"
 
         return new_aggs
 
@@ -361,9 +381,11 @@ class NotificationRuleData(object):
             pass
 
         # We hit the return above if we don't need to notify
-        if "aggs" in meta and "matched" not in meta:
+        meta["matched"] = self.format_matched(message)
+        if "aggs" in meta:
             meta["matched"] = self.format_aggs(meta["aggs"])
             print("MATCHED", meta["matched"])
+
         rule_notify(self.object, index, message, meta)
         self.store_match(index, message)
         await self.ingest_matches(index, message, meta, mode)