Thrift with SSL

1 files changed, 75 insertions, 3 deletions

diff --git a/module/remote/thriftbackend/Socket.py b/module/remote/thriftbackend/Socket.py
index 6ee850d07..cfb8b08c9 100644
--- a/module/remote/thriftbackend/Socket.py
+++ b/module/remote/thriftbackend/Socket.py
@@ -1,9 +1,27 @@
 # -*- coding: utf-8 -*-
 
+import sys
 import socket
+import errno
 
 from thrift.transport.TSocket import TSocket, TServerSocket, TTransportException
 
+class SecureSocketConnection:
+    def __init__(self, connection):
+        self.__dict__["connection"] = connection
+
+    def __getattr__(self, name):
+        return getattr(self.__dict__["connection"], name)
+
+    def __setattr__(self, name, value):
+        setattr(self.__dict__["connection"], name, value)
+
+    def shutdown(self, how=1):
+        self.__dict__["connection"].shutdown()
+
+    def accept(self):
+        connection, address = self.__dict__["connection"].accept()
+        return SecureSocketConnection(connection), address
 
 class Socket(TSocket):
     def __init__(self, host='localhost', port=7228, ssl=False):
@@ -11,21 +29,75 @@ class Socket(TSocket):
         self.ssl = ssl
 
     def open(self):
-        self.handle = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        if self.ssl:
+            SSL = __import__("OpenSSL", globals(), locals(), "SSL", -1).SSL
+            ctx = SSL.Context(SSL.SSLv23_METHOD)
+            self.handle = SecureSocketConnection(SSL.Connection(ctx, socket.socket(socket.AF_INET, socket.SOCK_STREAM)))
+        else:
+            self.handle = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+        #errno 104 connection reset
+
         self.handle.settimeout(self._timeout)
         self.handle.connect((self.host, self.port))
 
+    def read(self, sz):
+        try:
+            buff = self.handle.recv(sz)
+        except socket.error, e:
+            if (e.args[0] == errno.ECONNRESET and
+                (sys.platform == 'darwin' or sys.platform.startswith('freebsd'))):
+                # freebsd and Mach don't follow POSIX semantic of recv
+                # and fail with ECONNRESET if peer performed shutdown.
+                # See corresponding comment and code in TSocket::read()
+                # in lib/cpp/src/transport/TSocket.cpp.
+                self.close()
+                # Trigger the check to raise the END_OF_FILE exception below.
+                buff = ''
+            else:
+                raise
+        except Exception, e:
+            # SSL connection was closed
+            if e.args == (-1, 'Unexpected EOF'):
+                buff = ''
+            else:
+                raise
+            
+        if not len(buff):
+          raise TTransportException(type=TTransportException.END_OF_FILE, message='TSocket read 0 bytes')
+        return buff
+
 
 class ServerSocket(TServerSocket, Socket):
     def __init__(self, port=7228, host="0.0.0.0", key="", cert=""):
         self.host = host
         self.port = port
+        self.key = key
+        self.cert = cert
         self.handle = None
 
     def listen(self):
-        self.handle = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        if self.cert and self.key:
+            SSL = __import__("OpenSSL", globals(), locals(), "SSL", -1).SSL
+            ctx = SSL.Context(SSL.SSLv23_METHOD)
+            ctx.use_privatekey_file(self.key)
+            ctx.use_certificate_file(self.cert)
+
+            tmpConnection = SSL.Connection(ctx, socket.socket(socket.AF_INET, socket.SOCK_STREAM))
+            self.handle = SecureSocketConnection(tmpConnection)
+
+        else:
+            self.handle = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+
         self.handle.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         if hasattr(self.handle, 'set_timeout'):
           self.handle.set_timeout(None)
         self.handle.bind((self.host, self.port))
-        self.handle.listen(128)
\ No newline at end of file
+        self.handle.listen(128)
+
+    def accept(self):
+        client, addr = self.handle.accept()
+        result = Socket()
+        result.setHandle(client)
+        return result
\ No newline at end of file