feature: jinja2 templates for strings

Author: Lukas-Heindl

intelmq/bots/outputs/smtp_batch/output.py

@@ -26,6 +26,12 @@
 except ImportError:
     Envelope = None
 
+try:
+    import jinja2
+    jinja_env = jinja2.Environment()
+except ImportError:
+    jinja2 = None
+
 
 def hash_arbitrary(value: Any) -> bytes:
     value_bytes = None
@@ -48,6 +54,7 @@ class Mail:
 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}
     alternative_mails: Optional[str] = None
     bcc: Optional[list] = None
     email_from: str = ""
@@ -115,6 +122,10 @@ def set_cache(self):
     def init(self):
         if Envelope is None:
             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}:"
 
@@ -153,7 +164,7 @@ def cli_run(self):
 
         print("Preparing mail queue...")
         self.timeout = []
-        mails = [m for m in self.prepare_mails() if m]
+        mails = [m for m in self.prepare_mails() if m[0]]
 
         print("")
         if self.limit_results:
@@ -194,7 +205,9 @@ def cli_run(self):
                 count = 0
                 exit_code = 0
                 for mail in mails:
-                    succ = self.build_mail(mail, send=True)
+                    template_data = mail[1]
+                    mail = mail[0]
+                    succ = self.build_mail(mail, send=True, template_data=template_data)
                     if not succ:
                         exit_code = 1
                     else:
@@ -215,6 +228,7 @@ 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)
@@ -226,6 +240,7 @@ 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
@@ -244,7 +259,7 @@ 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, send=True, override_to=self.testing_to)])
+        count = sum([1 for mail in mails if self.build_mail(mail[0], send=True, override_to=self.testing_to, template_data=mail[1])])
         print(f"{count}× mail sent to: {self.testing_to}\n")
 
     def prepare_mails(self):
@@ -342,10 +357,10 @@ def prepare_mails(self):
 
             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)
-            yield mail
+            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 {})
 
-    def build_mail(self, mail, send=False, override_to=None):
+    def build_mail(self, mail, send=False, override_to=None, template_data:dict[str, Any]={}):
         """ creates a MIME message
         :param mail: Mail object
         :param send: True to send through SMTP, False for just printing the information
@@ -360,15 +375,30 @@ def build_mail(self, mail, send=False, override_to=None):
             intended_to = None
             email_to = mail.to
         email_from = self.email_from
+
         text = self.mail_contents
-        try:
-            subject = time.strftime(self.subject)
-        except ValueError:
-            subject = self.subject
-        try:
-            attachment_name = time.strftime(self.attachment_name)
-        except ValueError:
-            attachment_name = self.attachment_name
+        if jinja2 and self.templating and self.templating.get('body', False):
+            jinja_tmpl = jinja_env.from_string(text)
+            text = jinja_tmpl.render(current_time=datetime.datetime.now(), **template_data)
+
+        if jinja2 and self.templating and self.templating.get('subject', False):
+            jinja_tmpl = jinja_env.from_string(self.subject)
+            subject = jinja_tmpl.render(current_time=datetime.datetime.now(), **template_data)
+        else:
+            try:
+                subject = time.strftime(self.subject)
+            except ValueError:
+                subject = self.subject
+
+        if jinja2 and self.templating and self.templating.get('attachment', False):
+            jinja_tmpl = jinja_env.from_string(self.attachment_name)
+            attachment_name = jinja_tmpl.render(current_time=datetime.datetime.now(), **template_data)
+        else:
+            try:
+                attachment_name = time.strftime(self.attachment_name)
+            except ValueError:
+                attachment_name = self.attachment_name
+
         if intended_to:
             subject += f" (intended for {intended_to})"
         else: