Allow bounce header/footer to be strings

dldorsey · Jan 31, 2016 · 5cf416d · 5cf416d
1 parent d12a3d4
commit 5cf416d
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/slimta/bounce.py b/slimta/bounce.py
@@ -141,12 +141,18 @@ def __init__(self, envelope, reply, headers_only=False):
     def _check_custom_templates(cls):
         if cls.header_template != default_header_template and \
                 not isinstance(cls.header_template, BytesFormat):
-            fix_crlf = re.sub(br'\r?\n', br'\r\n', cls.header_template)
-            cls.header_template = BytesFormat(fix_crlf, mode='remove')
+            template_str = cls.header_template
+            if not isinstance(template_str, bytes):
+                template_str = template_str.encode('ascii')
+            template_str = re.sub(br'\r?\n', br'\r\n', template_str)
+            cls.header_template = BytesFormat(template_str, mode='remove')
         if cls.footer_template != default_footer_template and \
                 not isinstance(cls.footer_template, BytesFormat):
-            fix_crlf = re.sub(br'\r?\n', br'\r\n', cls.footer_template)
-            cls.footer_template = BytesFormat(fix_crlf, mode='remove')
+            template_str = cls.footer_template
+            if not isinstance(template_str, bytes):
+                template_str = template_str.encode('ascii')
+            template_str = re.sub(br'\r?\n', br'\r\n', template_str)
+            cls.footer_template = BytesFormat(template_str, mode='remove')
 
     def _get_substitution_table(self, envelope, reply, headers_only):
         rendered_rcpts = self.recipient_join.join(envelope.recipients)

diff --git a/test/test_slimta_bounce.py b/test/test_slimta_bounce.py
@@ -19,13 +19,13 @@ def test_bounce(self):
 """)
         reply = Reply('550', '5.0.0 Rejected')
 
-        Bounce.header_template = b"""\
+        Bounce.header_template = """\
 X-Reply-Code: {code}
 X-Reply-Message: {message}
 X-Orig-Sender: {sender}
 
 """
-        Bounce.footer_template = b"""\
+        Bounce.footer_template = """\
 
 EOM
 """
@@ -58,13 +58,13 @@ def test_bounce_headersonly(self):
 """)
         reply = Reply('550', '5.0.0 Rejected')
 
-        Bounce.header_template = b"""\
+        Bounce.header_template = """\
 X-Reply-Code: {code}
 X-Reply-Message: {message}
 X-Orig-Sender: {sender}
 
 """
-        Bounce.footer_template = b"""\
+        Bounce.footer_template = """\
 
 EOM
 """