Merge lp:~andrea.corbellini/beeseek/ssl-sockets into lp:beeseek/1.0

[Andrea Corbellini]

[Andrea Corbellini (andrea.corbellini)]

This branch enables sockets to receive and send data using SSL encryption.
It requires pyopenssl.

[Andrea Corbellini (andrea.corbellini)]

Approved after discussion with team members.

review: Approve

=== modified file 'beeseek/network/lowlevel.py'
--- beeseek/network/lowlevel.py	2009-02-17 17:06:44 +0000
+++ beeseek/network/lowlevel.py	2009-02-25 17:00:00 +0000
@@ -18,12 +18,20 @@
 from select import select
 from beeseek.interfaces import Interface, implements
 
+try:
+    from cStringIO import StringIO
+except ImportError:
+    from StringIO import StringIO
+
+
 __all__ = ('SocketError', 'IProtocol', 'BaseProtocol', 'ISocket',
            'BaseSocket', 'IPSocket', 'IPv6Socket', 'UNIXSocket',
-           'has_ipv6', 'has_unix')
+           'has_ipv6', 'has_unix', 'has_ssl', 'SSLContext')
 
 
 SocketError = _socket.error
+SSL_CLIENT = 0
+SSL_SERVER = 1
 
 class IProtocol(Interface):
     """Base interface for all network connectors.
@@ -142,6 +150,9 @@
 
     # Other socket methods
 
+    def fileno(self):
+        """Return the fileno."""
+
     def waitinput(self, timeout=None):
         """Wait for incoming data.
 
@@ -150,6 +161,16 @@
         :return: True if there is incoming data; else False.
         """
 
+    def enable_ssl(self, context=None, state=None):
+        """Read and write data using an SSL connection.
+
+        :param context: The SSLContext to use.
+        :param state: The connection state.
+        """
+
+    def disable_ssl(self):
+        """Read and write data using a standard connection."""
+
     # Close methods
 
     def close(self):
@@ -181,7 +202,8 @@
     high-level network objects.
     """
 
-    __slots__ = 'closed', '_sock', '_fileno', '_rfile', '_wfile'
+    __slots__ = ('closed', '_sock', '_fileno', '_rfile', '_wfile',
+                 '_rbuf', '_ssl')
     implements(IProtocol, isbase=True)
 
 
@@ -190,6 +212,7 @@
         self._fileno = fileno = self._sock.fileno()
         self._rfile = os.fdopen(fileno, 'r', 0)
         self._wfile = os.fdopen(fileno, 'w', bufsize)
+        self._ssl = None
         self.closed = False
 
     # Connect/bind methods
@@ -212,32 +235,111 @@
 
     def raw_recv(self, size=-1):
         if size < 0:
-            size = 5120
-        try:
-            data = os.read(self._fileno, size)
-        except IOError:
-            self.closed = True
-            return ''
+            size = 8192
+        if not self._ssl:
+            try:
+                data = os.read(self._fileno, size)
+            except IOError:
+                self.closed = True
+                return ''
+        else:
+            rbuf = self._rbuf
+            rbuf.seek(0, 0)
+            data = rbuf.read(size)
+            self._rbuf = StringIO(rbuf.read())
+            if data:
+                return data
+            data = self._ssl.recv(size)
         if not data:
             self.closed = True
         return data
 
     def raw_read(self, size=-1):
-        try:
-            data = self._rfile.read(size)
-        except IOError:
-            self.closed = True
-            return ''
+        if not self._ssl:
+            try:
+                data = self._rfile.read(size)
+            except IOError:
+                self.closed = True
+                return ''
+        else:
+            ssl = self._ssl
+            rbuf = self._rbuf
+            if size < 0:
+                self._rbuf = StringIO()
+                while True:
+                    data = ssl.recv(8192)
+                    if not data:
+                        break
+                    rbuf.write(data)
+            else:
+                rbuf.seek(0, 2)
+                buflen = rbuf.tell()
+                if buflen >= size:
+                    rbuf.seek(0, 0)
+                    data = rbuf.read(size)
+                    self._rbuf = StringIO(rbuf.read())
+                    return data
+                else:
+                    self._rbuf = StringIO()
+                    size -= buflen
+                    while size:
+                        data = ssl.recv(size)
+                        if not data:
+                            break
+                        rbuf.write(data)
+                        size -= len(data)
+            data = rbuf.getvalue()
         if not data:
             self.closed = True
         return data
 
     def raw_readline(self, size=-1):
-        try:
-            data = self._rfile.readline(size)
-        except IOError:
-            self.closed = True
-            return ''
+        if not self._ssl:
+            try:
+                data = self._rfile.readline(size)
+            except IOError:
+                self.closed = True
+                return ''
+        else:
+            ssl = self._ssl
+            rbuf = self._rbuf
+            rbuf.seek(0, 0)
+            if size < 0:
+                data = rbuf.readline()
+                if data.endswith('\n'):
+                    self._rbuf = StringIO()
+                    return data
+                while True:
+                    data = ssl.recv(8192)
+                    if not data:
+                        self._rbuf = StringIO()
+                        break
+                    newline = data.find('\n') + 1
+                    if newline > 0:
+                        rbuf.write(data[:newline])
+                        self._rbuf = StringIO(data[newline:])
+                        break
+                    rbuf.write(data)
+                data = rbuf.read()
+            else:
+                data = rbuf.readline(size)
+                size -= len(data)
+                if not size or data.endswith('\n'):
+                    self._rbuf = StringIO(rbuf.read())
+                    return data
+                self._rbuf = StringIO()
+                while size:
+                    data = ssl.recv(size)
+                    if not data:
+                        break
+                    newline = data.find('\n') + 1
+                    if newline > 0:
+                        rbuf.write(data[:newline])
+                        self._rbuf.write(data[newline:])
+                        break
+                    size -= len(data)
+                    rbuf.write(data)
+            data = rbuf.getvalue()
         if not data:
             self.closed = True
         return data
@@ -245,16 +347,22 @@
     # Send methods
 
     def raw_write(self, data):
-        try:
-            self._wfile.write(data)
-        except IOError:
-            self.closed = True
+        if not self._ssl:
+            try:
+                self._wfile.write(data)
+            except IOError:
+                self.closed = True
+        else:
+            self._ssl.sendall(data)
 
     def raw_writelines(self, data):
-        try:
-            self._wfile.writelines(data)
-        except IOError:
-            self.closed = True
+        if not self._ssl:
+            try:
+                self._wfile.writelines(data)
+            except IOError:
+                self.closed = True
+        else:
+            self._ssl.sendall(''.join(data))
 
     def flush(self):
         try:
@@ -268,12 +376,29 @@
         return self._fileno
 
     def waitinput(self, timeout=None):
-        return bool(select((self._rfile,), (), (), timeout))
+        return bool(select((self._rfile,), (), (), timeout)[0])
+
+    def enable_ssl(self, context=None, state=SSL_CLIENT):
+        if not context:
+            context = SSLContext()
+        self._ssl = ssl = SSL.Connection(context, self._sock)
+        if state == SSL_CLIENT:
+            ssl.set_connect_state()
+        elif state == SSL_SERVER:
+            ssl.set_accept_state()
+        self._rbuf = StringIO()
+
+    def disable_ssl(self):
+        self._ssl.shutdown()
+        self._ssl = None
+        del self._rbuf
 
     # Close methods
 
     def close(self):
         self.flush()
+        if self._ssl:
+            self.disable_ssl()
         try:
             self._sock.shutdown(_socket.SHUT_RDWR)
         except SocketError:
@@ -384,3 +509,20 @@
         family = _socket.AF_UNIX
 else:
     has_unix = False
+
+
+try:
+    from OpenSSL import SSL
+    from OpenSSL.SSL import ZeroReturnError
+except ImportError:
+    has_ssl = False
+else:
+    has_ssl = True
+    class SSLContext(object):
+
+        def __new__(cls, pemfile=None, certfile=None):
+            context = SSL.Context(SSL.SSLv23_METHOD)
+            if pemfile:
+                context.use_privatekey_file(pemfile)
+                context.use_certificate_file(certfile)
+            return context