Implement STARTTLS event; Fix #113

from circuits import Component, Debugger from circuits.net.sockets import TCPServer, write, starttls class TLSEchoServer(Component): def init(self, bind): self.transport = TCPServer(bind, certfile="/circuits/tests/net/cert.pem", ).register(self) def connect(self, sock, foo, bar): self.fire(write(sock, b'* OK [CAPABILITY STARTTLS] Dovecot ready.\r\n')) def read(self, sock, data): if data.strip().upper() == b". STARTTLS": yield self.fire(write(sock, b'. OK Begin TLS negotiation now.')) self.fire(starttls(sock)) return if data.strip() == b'. CAPABILITY': self.fire(write(sock, b'* CAPABILITY STARTTLS \r\n. OK Pre-login capabilities listed, post-login capabilities have more.\r\n')) else: self.fire(write(sock, b'Thanks for TLS encrypting! You wrote: %r' % (data,))) (TLSEchoServer(("0.0.0.0", 8000)) + Debugger()).run() $ openssl s_client -connect localhost:8000 -starttls imap
circuits · Jan 23, 2017 · d889942 · d889942
1 parent 25442fc
commit d889942
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 39 deletions.
diff --git a/circuits/net/events.py b/circuits/net/events.py
@@ -243,3 +243,15 @@ class closed(Event):
     .. note::
         This event is for Server components.
     """
+
+
+class starttls(Event):
+
+    """starttls Event
+
+    This event can be fired to upgrade the socket connection to a TLS
+    secured connection.
+
+    .. note::
+        This event is currently only available for Server Components.
+    """
diff --git a/circuits/net/sockets.py b/circuits/net/sockets.py
@@ -25,7 +25,7 @@
 
 from .events import (
     close, closed, connect, connected, disconnect, disconnected, error, read,
-    ready, unreachable, write,
+    ready, unreachable, write, starttls
 )
 
 try:
@@ -38,6 +38,8 @@
     import warnings
     warnings.warn("No SSL support available.")
     HAS_SSL = 0
+    CERT_NONE = None
+    PROTOCOL_SSLv23 = None
 
 
 BUFSIZE = 4096  # 4KB Buffer
@@ -407,18 +409,16 @@ def __init__(self, bind, secure=False, backlog=BACKLOG,
         self._poller = None
         self._buffers = defaultdict(deque)
 
-        self.secure = secure
+        self.__starttls = None
 
-        if self.secure:
-            try:
-                self.certfile = kwargs["certfile"]
-            except KeyError:
-                raise RuntimeError(
-                    "certfile must be specified for server-side operations")
-            self.keyfile = kwargs.get("keyfile", None)
-            self.cert_reqs = kwargs.get("cert_reqs", CERT_NONE)
-            self.ssl_version = kwargs.get("ssl_version", PROTOCOL_SSLv23)
-            self.ca_certs = kwargs.get("ca_certs", None)
+        self.secure = secure
+        self.certfile = kwargs.get("certfile")
+        self.keyfile = kwargs.get("keyfile", None)
+        self.cert_reqs = kwargs.get("cert_reqs", CERT_NONE)
+        self.ssl_version = kwargs.get("ssl_version", PROTOCOL_SSLv23)
+        self.ca_certs = kwargs.get("ca_certs", None)
+        if self.secure and not self.certfile:
+            raise RuntimeError("certfile must be specified for server-side operations")
 
     def parse_bind_parameter(self, bind_parameter):
         return parse_ipv4_parameter(bind_parameter)
@@ -564,20 +564,7 @@ def write(self, sock, data):
             self._poller.addWriter(self, sock)
         self._buffers[sock].append(data)
 
-    def _accept(self):  # noqa
-        # XXX: C901: This has a high McCacbe complexity score of 10.
-        # TODO: Refactor this!
-
-        def on_done(sock, host):
-            sock.setblocking(False)
-            self._poller.addReader(self, sock)
-            self._clients.append(sock)
-            self.fire(connect(sock, *host))
-
-        def on_error(sock, err):
-            self.fire(error(sock, err))
-            self._close(sock)
-
+    def _accept(self):
         try:
             newsock, host = self._sock.accept()
         except SocketError as e:
@@ -607,21 +594,47 @@ def on_error(sock, err):
                 raise
 
         if self.secure and HAS_SSL:
-            sslsock = ssl_socket(
-                newsock,
-                server_side=True,
-                keyfile=self.keyfile,
-                ca_certs=self.ca_certs,
-                certfile=self.certfile,
-                cert_reqs=self.cert_reqs,
-                ssl_version=self.ssl_version,
-                do_handshake_on_connect=False
-            )
-
-            for _ in do_handshake(sslsock, on_done, on_error, extra_args=(host,)):
+            for _ in self._do_handshake(newsock):
                 yield
         else:
-            on_done(newsock, host)
+            self._on_accept_done(newsock)
+
+    def _do_handshake(self, sock):
+        sslsock = ssl_socket(
+            sock,
+            server_side=True,
+            keyfile=self.keyfile,
+            ca_certs=self.ca_certs,
+            certfile=self.certfile,
+            cert_reqs=self.cert_reqs,
+            ssl_version=self.ssl_version,
+            do_handshake_on_connect=False
+        )
+
+        for _ in do_handshake(sslsock, self._on_accept_done, self._on_handshake_error):
+            yield _
+
+    def _on_accept_done(self, sock):
+        sock.setblocking(False)
+        self._poller.addReader(self, sock)
+        self._clients.append(sock)
+        self.fire(connect(sock, *sock.getpeername()))
+
+    def _on_handshake_error(self, sock, err):
+        self.fire(error(sock, err))
+        self._close(sock)
+
+    @handler('starttls')
+    def starttls(self, sock):
+        if not HAS_SSL:
+            return
+        if self.__starttls:
+            raise RuntimeError('starttls started twice.')
+        self.__starttls = True
+        self._poller.removeReader(sock)
+        self._clients.remove(sock)
+        for _ in self._do_handshake(sock):
+            yield
 
     @handler("_disconnect", priority=1)
     def _on_disconnect(self, sock):