Update test_smtpsmuggling.py (#416)

---------

Co-authored-by: Sam Bull <git@sambull.org>

aiosmtpd/tests/test_smtpsmuggling.py

@@ -3,53 +3,65 @@
 
 """Test SMTP smuggling."""
 
-from email.mime.text import MIMEText
-from smtplib import SMTP, SMTP_SSL
-from typing import Generator, Union
-
-import pytest
 import smtplib
+import re
 
-from aiosmtpd.controller import Controller
 from aiosmtpd.testing.helpers import ReceivingHandler
 from aiosmtpd.testing.statuscodes import SMTP_STATUS_CODES as S
 
-from aiosmtpd.smtp import SMTP as Server
-from aiosmtpd.smtp import Session as ServerSession
-from aiosmtpd.smtp import Envelope
-
-from .conftest import Global, controller_data, handler_data
+from .conftest import handler_data
 
-from aiosmtpd.testing.helpers import (
-    ReceivingHandler
-)
 
 def new_data(self, msg):
-        self.putcmd("data")
+    self.putcmd("data")
+    (code, repl) = self.getreply()
+    if self.debuglevel > 0:
+        self._print_debug('data:', (code, repl))
+    if code != 354:
+        raise smtplib.SMTPDataError(code, repl)
+    else:
+        q = msg
+        self.send(q)
+        (code, msg) = self.getreply()
+        if self.debuglevel > 0:
+            self._print_debug('data:', (code, msg))
+        return (code, msg)
 
-        (code, repl) = self.getreply()
+
+def orig_data(self, msg):
+    self.putcmd("data")
+    (code, repl) = self.getreply()
+    if self.debuglevel > 0:
+        self._print_debug('data:', (code, repl))
+    if code != 354:
+        raise smtplib.SMTPDataError(code, repl)
+    else:
+        if isinstance(msg, str):
+            msg = _fix_eols(msg).encode('ascii')
+        q = _quote_periods(msg)
+        if q[-2:] != smtplib.bCRLF:
+            q = q + smtplib.bCRLF
+
+        q = q + b"." + smtplib.bCRLF
+        self.send(q)
+        (code, msg) = self.getreply()
         if self.debuglevel > 0:
-            self._print_debug('data:', (code, repl))
-        if code != 354:
-            raise SMTPDataError(code, repl)
-        else:
-            ##### Patching input encoding so we can send raw messages
-            #if isinstance(msg, str):
-            #    msg = smtplib._fix_eols(msg).encode('ascii')
-            #q = smtplib._quote_periods(msg)
-            #if q[-2:] != smtplib.bCRLF:
-            #    q = q + smtplib.bCRLF
-            #q = q + b"." + smtplib.bCRLF
-            q = msg
-            self.send(q)
-            (code, msg) = self.getreply()
-            if self.debuglevel > 0:
-                self._print_debug('data:', (code, msg))
-            return (code, msg)
+            self._print_debug('data:', (code, msg))
+        return (code, msg)
+
+
+def _fix_eols(data):
+    return re.sub(r'(?:\r\n|\n|\r(?!\n))', smtplib.CRLF, data)
+
+
+def _quote_periods(bindata):
+    return re.sub(br'(?m)^\.', b'..', bindata)
+
 
 def return_unchanged(data):
     return data
 
+
 class TestSmuggling:
     @handler_data(class_=ReceivingHandler)
     def test_smtp_smuggling(self, plain_controller, client):
@@ -74,6 +86,10 @@ def test_smtp_smuggling(self, plain_controller, client):
 NO SMUGGLING
 \r\n.\r\n\
 """
-        results = client.sendmail(sender, recipients, message_data)
+        client.sendmail(sender, recipients, message_data)
         client.quit()
+        smtplib._fix_eols = _fix_eols
+        smtplib._quote_periods = _quote_periods
+        smtplib.SMTP.data = orig_data
+
         assert b"NO SMUGGLING" in handler.box[0].content