fixup unit test failure

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, Generator, Optional
+from typing import Any, Iterable, Optional, Dict, List
 import zipfile
 from base64 import b64decode
 from collections import OrderedDict
@@ -50,13 +50,13 @@ class Mail:
     to: str
     path: str
     count: int
-    template_data: dict[str,Any]
+    template_data: Dict[str, Any]
 
 
 class SMTPBatchOutputBot(Bot):
     # configurable parameters
     additional_grouping_keys: Optional[list] = []  # refers to the event directly
-    templating: Optional[dict[str, bool]] = {'subject': False, 'body': False, 'attachment': False}
+    templating: Optional[Dict[str, bool]] = {'subject': False, 'body': False, 'attachment': False}
     alternative_mails: Optional[str] = None
     bcc: Optional[list] = None
     email_from: str = ""
@@ -126,8 +126,6 @@ def init(self):
             raise MissingDependencyError('envelope', '>=2.0.0')
         if jinja2 is None:
             self.logger.warning("No jinja2 installed. Thus, the templating is deactivated.")
-        if self.additional_grouping_keys and len(self.additional_grouping_keys) > 0:
-            self.additional_grouping_keys_trans = {self.fieldnames_translation.get(k, k) for k in self.additional_grouping_keys}
         self.set_cache()
         self.key = f"{self._Bot__bot_id}:"
 
@@ -251,7 +249,7 @@ def set_tester(self, force=True):
         print("\nWhat e-mail should I use?")
         self.testing_to = input()
 
-    def send_mails_to_tester(self, mails):
+    def send_mails_to_tester(self, mails: List[Mail]):
         """
             These mails are going to tester's address. Then prints out their count.
         :param mails: list
@@ -260,7 +258,7 @@ def send_mails_to_tester(self, mails):
         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) -> Generator[Mail]:
+    def prepare_mails(self) -> Iterable[Mail]:
         """ Generates Mail objects """
 
         for mail_record in self.cache.redis.keys(f"{self.key}*")[slice(self.limit_results)]:
@@ -332,14 +330,16 @@ def prepare_mails(self) -> Generator[Mail]:
 
             # collect all data which must be the same for all events of the
             # bucket and thus can be used for templating
-            template_data = {}
-            # only collect if templating is enabled (save the memory otherwise)
+            template_keys = ['source.abuse_contact']
+            # 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]
-                }
+                template_keys.extend(self.additional_grouping_keys)
+
+            template_data = {
+                k.replace(".", "_"): lines[0][k]
+                for k in template_keys
+                if k in lines[0]
+            }
 
             email_to = template_data["source_abuse_contact"]
             filename = f'{time.strftime("%y%m%d")}_{count}_events'
@@ -361,7 +361,7 @@ def prepare_mails(self) -> Generator[Mail]:
             self.build_mail(mail, send=False)
             yield mail
 
-    def build_mail(self, mail, send=False, override_to=None):
+    def build_mail(self, mail: 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
@@ -377,6 +377,8 @@ def build_mail(self, mail, send=False, override_to=None):
             email_to = mail.to
         email_from = self.email_from
 
+        template_data = mail.template_data
+
         text = self.mail_contents
         if jinja2 and self.templating and self.templating.get('body', False):
             jinja_tmpl = jinja_env.from_string(text)

intelmq/tests/bots/outputs/smtp_batch/test_output.py

@@ -11,19 +11,21 @@
 
 import intelmq.lib.test as test
 from intelmq.bots.outputs.smtp_batch.output import SMTPBatchOutputBot
-from intelmq.lib.cache import Cache
 from intelmq.lib.exceptions import MissingDependencyError
+from hashlib import sha256
 
 BOT_ID = "test-bot"
 IDENTITY1 = '[email protected]'
-KEY1 = f"{BOT_ID}:{IDENTITY1}".encode()
+IDENTITY1_HASH = sha256(sha256(IDENTITY1.encode()).digest()).hexdigest()
+KEY1 = f"{BOT_ID}:{IDENTITY1_HASH}".encode()
 EVENT1 = {'__type': 'Event',
           'source.ip': '127.0.0.1',
           'source.url': 'http://example.com/',
           'source.abuse_contact': IDENTITY1
           }
 IDENTITY2 = '[email protected]'
-KEY2 = f"{BOT_ID}:{IDENTITY2}".encode()
+IDENTITY2_HASH = sha256(sha256(IDENTITY2.encode()).digest()).hexdigest()
+KEY2 = f"{BOT_ID}:{IDENTITY2_HASH}".encode()
 EVENT2 = {'__type': 'Event',
           'source.ip': '127.0.0.2',
           'source.url': 'http://example2.com/',
@@ -67,6 +69,7 @@ def compare_envelope(self, envelope: Envelope, subject, message, from_, to):
         self.assertEqual(message, envelope.message())
         self.assertEqual(from_, envelope.from_())
         self.assertEqual(to, envelope.to())
+        # TODO template_data
 
     def send_message(self):
         def _(envelope):