Add total hits to output

core/db/elastic.py

@@ -300,8 +300,8 @@ class ElasticsearchBackend(StorageBackend):
                 search_query,
                 index=index,
             )
-            self.log.debug(f"Running scheduled query on {index}: {search_query}")
-            self.log.debug(f"Response from scheduled query: {response}")
+            # self.log.debug(f"Running scheduled query on {index}: {search_query}")
+            # self.log.debug(f"Response from scheduled query: {response}")
             if isinstance(response, Exception):
                 error = response.info["error"]["root_cause"][0]["reason"]
                 self.log.error(f"Error running scheduled search: {error}")
@@ -310,7 +310,7 @@ class ElasticsearchBackend(StorageBackend):
                 # No results, skip
                 continue
             meta, response = self.parse(response, meta=True)
-            print("Parsed response", response)
+            # print("Parsed response", response)
             if "message" in response:
                 self.log.error(f"Error running scheduled search: {response['message']}")
                 continue

core/db/processing.py

@@ -125,6 +125,8 @@ def parse_results(results, meta=None):
             for field in ["avg_sentiment"]:  # Add other number fields here
                 if field in results["aggregations"]:
                     meta["aggs"][field] = results["aggregations"][field]
+        total_hits = results["hits"]["total"]["value"]
+        meta["total_hits"] = total_hits
         return (meta, results_parsed)
 
     return results_parsed

core/lib/rules.py

@@ -41,11 +41,13 @@ def format_ntfy(**kwargs):
         index: The index the rule matched on, can be None
         message: The message to send, can be None
         matched: The matched fields, can be None
+        total_hits: The total number of matches, optional
     """
     rule = kwargs.get("rule")
     index = kwargs.get("index")
     message = kwargs.get("message")
     matched = kwargs.get("matched")
+    total_hits = kwargs.get("total_hits", 0)
     if message:
         # Dump the message in YAML for readability
         messages_formatted = ""
@@ -66,6 +68,7 @@ def format_ntfy(**kwargs):
         matched = ""
 
     notify_message = f"{rule.name} on {index}: {matched}\n{messages_formatted}"
+    notify_message += f"\nTotal hits: {total_hits}"
     notify_message = notify_message.encode("utf-8", "replace")
 
     return notify_message
@@ -82,6 +85,7 @@ def format_webhook(**kwargs):
         index: The index the rule matched on, can be None
         message: The message to send, can be None, but will be sent as None
         matched: The matched fields, can be None, but will be sent as None
+        total_hits: The total number of matches, optional
         notification_settings: The notification settings, must be specified
             priority: The priority of the message, optional
             topic: The topic of the message, optional
@@ -90,11 +94,13 @@ def format_webhook(**kwargs):
     index = kwargs.get("index")
     message = kwargs.get("message")
     matched = kwargs.get("matched")
+    total_hits = kwargs.get("total_hits", 0)
     notification_settings = kwargs.get("notification_settings")
     notify_message = {
         "rule_id": rule.id,
         "rule_name": rule.name,
         "match": matched,
+        "total_hits": total_hits,
         "index": index,
         "data": message,
     }
@@ -141,7 +147,10 @@ def rule_notify(rule, index, message, meta=None):
         "notification_settings": notification_settings,
     }
     if meta:
-        cast["matched"] = meta["aggs"]
+        if "matched" in meta:
+            cast["matched"] = meta["matched"]
+        if "total_hits" in meta:
+            cast["total_hits"] = meta["total_hits"]
 
     if rule.service == "ntfy":
         cast["msg"] = format_ntfy(**cast)
@@ -263,7 +272,7 @@ class NotificationRuleData(object):
         log.debug(f"Rule matched: {index} - current match: {current_match}")
         if current_match is False:
             # Matched now, but not before
-            meta["aggs"] = self.format_aggs(meta["aggs"])
+            meta["matched"] = self.format_aggs(meta["aggs"])
             rule_notify(self.object, index, message, meta)
             self.store_match(index, True)
 

core/management/commands/processing.py

@@ -78,7 +78,8 @@ def process_rules(data):
                         matched[field] = message[field]
                 # Subtract 2, 1 for source and 1 for index
                 if matched_field_number == rule_field_length - 2:
-                    rule_notify(rule, index, message, matched)
+                    meta = {"matched": matched, "total_hits": 1}
+                    rule_notify(rule, index, message, meta=meta)
 
 
 class Command(BaseCommand):