Repository: aio-libs/aiosmtpd
Pull request #64 (2017-04-07 20:40:59)

Rename handle_tls_handshake() to handle_STARTTLS
This builds on, but changes #53 and should replace that branch.


Patch:
diff --git a/aiosmtpd/docs/NEWS.rst b/aiosmtpd/docs/NEWS.rst
index 2452ffe..84f57ea 100644
--- a/aiosmtpd/docs/NEWS.rst
+++ b/aiosmtpd/docs/NEWS.rst
@@ -21,9 +21,8 @@
 * Controller objects now have an optional timeout argument used to wait for
   the server to become ready.  This can also be overridden with the
   environment variable ``AIOSMTPD_CONTROLLER_TIMEOUT``. (Closes #35)
-* Handlers can define a ``handle_tls_handshake()`` method, which takes a
-  session object, and is called if SSL is enabled during the making of the
-  connection.  (Closes #48)
+* Handlers can define a ``handle_STARTTLS()`` method which is called if SSL is
+  enabled during the making of the connection. (Closes #48)
 * Better Windows compatibility.
 * Better Python 3.4 compatibility.
 * Use ``flufl.testing`` package for nose2 and flake8 plugins.
diff --git a/aiosmtpd/docs/smtp.rst b/aiosmtpd/docs/smtp.rst
index a7626b6..a4eb0aa 100644
--- a/aiosmtpd/docs/smtp.rst
+++ b/aiosmtpd/docs/smtp.rst
@@ -161,7 +161,7 @@ The following hooks are currently defined:
 ``handle_RSET(server, session, envelope)``
     Called during ``RSET``.
 
-``handle_DATA(session, envelope)``
+``handle_DATA(server, session, envelope)``
     Called during ``DATA`` after the entire message (`"SMTP content"
     <https://tools.ietf.org/html/rfc5321#section-2.3.9>`_ as described in
     RFC 5321) has been received.  The content is available on the ``envelope``
@@ -176,13 +176,13 @@ The following hooks are currently defined:
     content.
 
 In addition to the SMTP command hooks, the following hooks can also be
-implemented by handlers.  They have a different signature and don't need to
-return any status.  Both methods are called synchronously.
+implemented by handlers.  These have a different APIs, and are called
+synchronously.
 
-``handle_tls_handshake(session)``
-    If implemented, and if SSL is supported, this handler method gets called
+``handle_STARTTLS(server, session, envelope)``
+    If implemented, and if SSL is supported, this method gets called
     during the TLS handshake phase of ``connection_made()``.  It should return
-    a boolean which specifies whether the handshake failed or not.
+    True if the handshake succeeded, and False otherwise.
 
 ``handle_exception(error)``
     If implemented, this method is called when any error occurs during the
diff --git a/aiosmtpd/smtp.py b/aiosmtpd/smtp.py
index cf539c7..54d6157 100644
--- a/aiosmtpd/smtp.py
+++ b/aiosmtpd/smtp.py
@@ -92,7 +92,7 @@ class SMTP(asyncio.StreamReaderProtocol):
             self.tls_context.check_hostname = False
             self.tls_context.verify_mode = ssl.CERT_NONE
         self.require_starttls = tls_context and require_starttls
-        self._tls_handshake_failed = False
+        self._tls_handshake_okay = True
         self._tls_protocol = None
         self.session = None
         self.envelope = None
@@ -104,6 +104,14 @@ class SMTP(asyncio.StreamReaderProtocol):
     def _create_envelope(self):
         return Envelope()
 
+    @asyncio.coroutine
+    def _call_handler_hook(self, command, *args):
+        hook = getattr(self.event_handler, 'handle_' + command, None)
+        if hook is None:
+            return MISSING
+        status = yield from hook(self, self.session, self.envelope, *args)
+        return status
+
     @property
     def max_command_size_limit(self):
         try:
@@ -127,11 +135,12 @@ class SMTP(asyncio.StreamReaderProtocol):
             # Do SSL certificate checking as rfc3207 part 4.1 says.
             # Why _extra is protected attribute?
             self.session.ssl = self._tls_protocol._extra
-            if hasattr(self.event_handler, 'handle_tls_handshake'):
-                auth = self.event_handler.handle_tls_handshake(self.session)
-                self._tls_handshake_failed = not auth
+            handler = getattr(self.event_handler, 'handle_STARTTLS', None)
+            if handler is None:
+                self._tls_handshake_okay = True
             else:
-                self._tls_handshake_failed = False
+                self._tls_handshake_okay = handler(
+                    self, self.session, self.envelope)
             self._over_ssl = True
         else:
             super().connection_made(transport)
@@ -181,14 +190,6 @@ class SMTP(asyncio.StreamReaderProtocol):
         if hasattr(self.event_handler, 'handle_exception'):
             yield from self.event_handler.handle_exception(error)
 
-    @asyncio.coroutine
-    def _call_handler_hook(self, command, *args):
-        hook = getattr(self.event_handler, 'handle_' + command, None)
-        if hook is None:
-            return MISSING
-        status = yield from hook(self, self.session, self.envelope, *args)
-        return status
-
     @asyncio.coroutine
     def _handle_client(self):
         log.info('%r handling connection', self.session.peer)
@@ -217,7 +218,7 @@ class SMTP(asyncio.StreamReaderProtocol):
                 if len(line) > max_sz:
                     yield from self.push('500 Error: line too long')
                     continue
-                if (self._tls_handshake_failed
+                if (not self._tls_handshake_okay
                         and command != 'QUIT'):             # pragma: nossl
                     yield from self.push(
                         '554 Command refused due to lack of security')


Test patch:
diff --git a/aiosmtpd/tests/test_starttls.py b/aiosmtpd/tests/test_starttls.py
index 8fee782..f248c2e 100644
--- a/aiosmtpd/tests/test_starttls.py
+++ b/aiosmtpd/tests/test_starttls.py
@@ -24,13 +24,13 @@ class Controller(BaseController):
 
 
 class ReceivingHandler:
-    box = None
+    def __init__(self):
+        self.box = []
 
     @asyncio.coroutine
-    def process_message(self, *args, **kws):
-        if not self.box:
-            self.box = []
-        self.box.append(args)
+    def handle_DATA(self, server, session, envelope):
+        self.box.append(envelope)
+        return '250 OK'
 
 
 def get_tls_context():
@@ -60,8 +60,7 @@ class TLSController(Controller):
 
 
 class HandshakeFailingHandler:
-    @staticmethod
-    def handle_tls_handshake(session):
+    def handle_STARTTLS(self, server, session, envelope):
         return False