move template_data from tuple to Mail

Author: Lukas-Heindl

intelmq/bots/outputs/smtp_batch/output.py

@@ -8,7 +8,7 @@
 import sys
 from tempfile import NamedTemporaryFile
 import time
-from typing import Any, Optional
+from typing import Any, Generator, Optional
 import zipfile
 from base64 import b64decode
 from collections import OrderedDict
@@ -49,6 +49,7 @@ class Mail:
     to: str
     path: str
     count: int
+    template_data: dict[str,Any]
 
 
 class SMTPBatchOutputBot(Bot):
@@ -164,7 +165,7 @@ def cli_run(self):
 
         print("Preparing mail queue...")
         self.timeout = []
-        mails = [m for m in self.prepare_mails() if m[0]]
+        mails = [m for m in self.prepare_mails() if m]
 
         print("")
         if self.limit_results:
@@ -205,9 +206,7 @@ def cli_run(self):
                 count = 0
                 exit_code = 0
                 for mail in mails:
-                    template_data = mail[1]
-                    mail = mail[0]
-                    succ = self.build_mail(mail, send=True, template_data=template_data)
+                    succ = self.build_mail(mail, send=True)
                     if not succ:
                         exit_code = 1
                     else:
@@ -228,7 +227,6 @@ def cli_run(self):
                 sys.exit(exit_code)
             elif i == "clear":
                 for mail in mails:
-                    mail = mail[0]
                     self.cache.redis.delete(mail.key)
                 print("Queue cleared.")
                 sys.exit(0)
@@ -240,7 +238,6 @@ def cli_run(self):
                 self.send_mails_to_tester(mails)
             else:
                 for mail in mails:
-                    mail = mail[0]
                     if mail.to == i:
                         self.send_mails_to_tester([mail])
                         break
@@ -259,10 +256,10 @@ def send_mails_to_tester(self, mails):
         :param mails: list
         """
         self.set_tester(False)
-        count = sum([1 for mail in mails if self.build_mail(mail[0], send=True, override_to=self.testing_to, template_data=mail[1])])
+        count = sum([1 for mail in mails if self.build_mail(mail, send=True, override_to=self.testing_to)])
         print(f"{count}× mail sent to: {self.testing_to}\n")
 
-    def prepare_mails(self):
+    def prepare_mails(self) -> Generator[Mail]:
         """ Generates Mail objects """
 
         for mail_record in self.cache.redis.keys(f"{self.key}*")[slice(self.limit_results)]:
@@ -334,11 +331,14 @@ def prepare_mails(self):
 
             # 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]
-            }
+            template_data = {}
+            # only collect if templating is enabled (save the memory otherwise)
+            if jinja2 and self.templating and any(self.templating.values()):
+                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'
@@ -355,12 +355,12 @@ def prepare_mails(self):
                 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)
+            mail = Mail(mail_record, email_to, path, count, template_data)
             # build_mail only used to output metadata of the mail -> send=False -> return None
-            self.build_mail(mail, send=False, template_data=template_data)
-            yield (mail, template_data if jinja2 and self.templating and any(self.templating.values()) else {})
+            self.build_mail(mail, send=False)
+            yield mail
 
-    def build_mail(self, mail, send=False, override_to=None, template_data:dict[str, Any]={}):
+    def build_mail(self, mail, send=False, override_to=None):
         """ creates a MIME message
         :param mail: Mail object
         :param send: True to send through SMTP, False for just printing the information