smtp_batch: add option to group-by more keys than the dst mail-address

Author: Lukas-Heindl

intelmq/bots/outputs/smtp_batch/output.py

@@ -13,6 +13,7 @@
 from base64 import b64decode
 from collections import OrderedDict
 from io import StringIO
+from hashlib import sha256
 
 from redis.exceptions import TimeoutError
 
@@ -26,6 +27,16 @@
     Envelope = None
 
 
+def hash_arbitrary(value: Any) -> bytes:
+    value_bytes = None
+    if isinstance(value, str):
+        value_bytes = value.encode("utf-8")
+    elif isinstance(value, int):
+        value_bytes = bytes(value)
+    else:
+        value_bytes = json.dumps(value, sort_keys=True).encode("utf-8")
+    return sha256(value_bytes).digest()
+
 @dataclass
 class Mail:
     key: str
@@ -36,6 +47,7 @@ class Mail:
 
 class SMTPBatchOutputBot(Bot):
     # configurable parameters
+    additional_grouping_keys: Optional[list] = []  # refers to the event directly
     alternative_mails: Optional[str] = None
     bcc: Optional[list] = None
     email_from: str = ""
@@ -75,7 +87,20 @@ def process(self):
         if "source.abuse_contact" in message:
             field = message["source.abuse_contact"]
             for mail in (field if isinstance(field, list) else [field]):
-                self.cache.redis.rpush(f"{self.key}{mail}", message.to_json())
+                # - Each event goes into one bucket (equivalent to group-by)
+                # - The id of each bucket is calcuated by hashing all the keys that should be grouped for
+                # - Hashing ensures the redis-key does not grow indefinitely.
+                # - In order to avoid collisions, each value is hashed before
+                #   appending to the input for the redis-key-hash
+                #   (could also be solved by special separator which would need
+                #    to be escaped or prepending the length of the value).
+                h = sha256()
+                h.update(sha256(mail.encode("utf-8")).digest())
+                for i in self.additional_grouping_keys:
+                    if i not in message:
+                        continue
+                    h.update(hash_arbitrary(message[i]))
+                self.cache.redis.rpush(f"{self.key}{h.hexdigest()}", message.to_json())
 
         self.acknowledge_message()
 
@@ -254,7 +279,11 @@ def prepare_mails(self):
             # TODO: worthy to generate on the fly https://github.com/certtools/intelmq/pull/2253#discussion_r1172779620
             fieldnames = set()
             rows_output = []
+            src_abuse_contact = None
             for row in lines:
+                # obtain this field only once as it is the same for all lines here
+                if not src_abuse_contact:
+                    src_abuse_contact = row["source.abuse_contact"]
                 try:
                     if threshold and row["time.observation"][:19] < threshold.isoformat()[:19]:
                         continue
@@ -283,29 +312,38 @@ def prepare_mails(self):
             dict_writer.writerow(dict(zip(ordered_fieldnames, ordered_fieldnames)))
             dict_writer.writerows(rows_output)
 
-            email_to = str(mail_record[len(self.key):], encoding="utf-8")
             count = len(rows_output)
-            if not count:
-                path = None
-            else:
-                filename = f'{time.strftime("%y%m%d")}_{count}_events'
-                path = NamedTemporaryFile().name
-
-                with zipfile.ZipFile(path, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
-                    try:
-                        zf.writestr(filename + ".csv", output.getvalue())
-                    except Exception:
-                        self.logger.error("Error: Cannot zip mail: %r", mail_record)
-                        continue
+            if not count or count == 0:
+                # send no mail if no events are present
+                continue
+
+            # collect all data which must be the same for all events of the
+            # bucket and thus can be used for templating
+            template_data = {
+                k.replace(".", "_"): lines[0][k]
+                for k in ["source.abuse_contact"] + self.additional_grouping_keys
+                if k in rows_output[0]
+            }
+
+            email_to = template_data["source_abuse_contact"]
+            filename = f'{time.strftime("%y%m%d")}_{count}_events'
+            path = NamedTemporaryFile().name
+
+            with zipfile.ZipFile(path, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
+                try:
+                    zf.writestr(filename + ".csv", output.getvalue())
+                except Exception:
+                    self.logger.error("Error: Cannot zip mail: %r", mail_record)
+                    continue
 
-                if email_to in self.alternative_mail:
-                    print(f"Alternative: instead of {email_to} we use {self.alternative_mail[email_to]}")
-                    email_to = self.alternative_mail[email_to]
+            if email_to in self.alternative_mail:
+                print(f"Alternative: instead of {email_to} we use {self.alternative_mail[email_to]}")
+                email_to = self.alternative_mail[email_to]
 
             mail = Mail(mail_record, email_to, path, count)
+            # build_mail only used to output metadata of the mail -> send=False -> return None
             self.build_mail(mail, send=False)
-            if count:
-                yield mail
+            yield mail
 
     def build_mail(self, mail, send=False, override_to=None):
         """ creates a MIME message