Merge WebSocketServer and SSLWebSocketServer into WebSocketServer

setup.py

@@ -3,7 +3,7 @@
 
 setup(
     name='simple-websocket-server',
-    version='0.3.1',
+    version='0.4.0',
     author='Sergei Pikhovkin',
     author_email='[email protected]',
     packages=find_packages(exclude=['tests*']),

simple_websocket_server/__init__.py

@@ -21,8 +21,7 @@
 
 __all__ = [
     'WebSocket',
-    'WebSocketServer',
-    'SSLWebSocketServer'
+    'WebSocketServer'
 ]
 
 
@@ -553,7 +552,9 @@ def _parse_message(self, byte):  # pylint: disable=too-many-branches, too-many-s
 class WebSocketServer(object):
     request_queue_size = 5
 
-    def __init__(self, host, port, websocketclass, select_interval=0.1):
+    # pylint: disable=too-many-arguments
+    def __init__(self, host, port, websocketclass, certfile=None, keyfile=None,
+                 ssl_version=ssl.PROTOCOL_TLSv1, select_interval=0.1):
         self.websocketclass = websocketclass
         self.serversocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.serversocket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
@@ -563,11 +564,24 @@ def __init__(self, host, port, websocketclass, select_interval=0.1):
         self.connections = {}
         self.listeners = [self.serversocket]
 
+        self._using_ssl = bool(certfile and keyfile)
+
+        if self._using_ssl:
+            self.context = ssl.SSLContext(ssl_version)
+            self.context.load_cert_chain(certfile, keyfile)
+
     def _decorate_socket(self, sock):  # pylint: disable=no-self-use
+        if self._using_ssl:
+            return self.context.wrap_socket(sock, server_side=True)
+
         return sock
 
     def _construct_websocket(self, sock, address):
-        return self.websocketclass(self, sock, address)
+        ws = self.websocketclass(self, sock, address)
+        if self._using_ssl:
+            ws.usingssl = True
+
+        return ws
 
     def close(self):
         self.serversocket.close()
@@ -656,22 +670,3 @@ def handle_request(self):  # pylint: disable=too-many-branches, too-many-stateme
     def serve_forever(self):
         while True:
             self.handle_request()
-
-
-class SSLWebSocketServer(WebSocketServer):
-    # pylint: disable=too-many-arguments
-    def __init__(self, host, port, websocketclass, certfile, keyfile,
-                 version=ssl.PROTOCOL_TLSv1, select_interval=0.1):
-        WebSocketServer.__init__(self, host, port, websocketclass, select_interval)
-
-        self.context = ssl.SSLContext(version)
-        self.context.load_cert_chain(certfile, keyfile)
-
-    def _decorate_socket(self, sock):
-        sslsock = self.context.wrap_socket(sock, server_side=True)
-        return sslsock
-
-    def _construct_websocket(self, sock, address):
-        ws = self.websocketclass(self, sock, address)
-        ws.usingssl = True
-        return ws

tests/example_server.py

@@ -5,7 +5,7 @@
 import ssl
 from optparse import OptionParser
 
-from simple_websocket_server import WebSocket, WebSocketServer, SSLWebSocketServer
+from simple_websocket_server import WebSocket, WebSocketServer
 
 
 class SimpleEcho(WebSocket):
@@ -51,10 +51,11 @@ def handle_close(self):
     if options.example == 'chat':
         cls = SimpleChat
 
+    sslopts = {}
     if options.ssl == 1:
-        server = SSLWebSocketServer(options.host, options.port, cls, options.cert, options.key, version=options.ver)
-    else:
-        server = WebSocketServer(options.host, options.port, cls)
+        sslopts = dict(certfile=options.cert, keyfile=options.key, ssl_version=options.ver)
+
+    server = WebSocketServer(options.host, options.port, cls, **sslopts)
 
     def close_sig_handler(signal, frame):
         server.close()