Merge pull request Lawouach#33 from patrickod/master

Lawouach · Lawouach · commit b1bb563df66f · 2012-03-20T14:47:08.000-07:00
Fix SSL support in TornadoWebSocketClient and add missing import
diff --git a/ws4py/client/tornadoclient.py b/ws4py/client/tornadoclient.py
@@ -1,17 +1,26 @@
 # -*- coding: utf-8 -*-
 import socket
 from urlparse import urlsplit
+import ssl
 
 from tornado import iostream, escape
 from ws4py.client import WebSocketBaseClient
+from ws4py.exc import HandshakeError
 
 __all__ = ['TornadoWebSocketClient']
 
 class TornadoWebSocketClient(WebSocketBaseClient):
     def __init__(self, url, protocols=None, extensions=None, io_loop=None):
         WebSocketBaseClient.__init__(self, url, protocols, extensions)
-        self.io = iostream.IOStream(self.sock, io_loop)
+        parts = urlsplit(self.url)
+        if parts.scheme == "wss":
+            self.sock = ssl.wrap_socket(self.sock,
+                    do_handshake_on_connect=False)
+            self.io = iostream.SSLIOStream(self.sock, io_loop)
+        else:
+            self.io = iostream.IOStream(self.sock, io_loop)
         self.sender = self.io.write
+        self.io_loop = io_loop
 
     def connect(self):
         parts = urlsplit(self.url)
@@ -29,11 +38,11 @@ def __send_handshake(self):
         self.io.set_close_callback(self.__connection_closed)
         self.io.write(escape.utf8(self.handshake_request),
                       self.__handshake_sent)
-    
+
     def __connection_closed(self, *args, **kwargs):
         self.server_terminated = True
         self.closed(1006, 'Connection closed during handshake')
-    
+
     def __handshake_sent(self):
         self.io.read_until("\r\n\r\n", self.__handshake_completed)
 
@@ -46,7 +55,7 @@ def __handshake_completed(self, data):
         except HandshakeError:
             self.close_connection()
             raise
-        
+
         self.opened()
         self.io.set_close_callback(self.__stream_closed)
         self.io.read_bytes(self.reading_buffer_size, self.__fetch_more)
@@ -64,7 +73,7 @@ def __fetch_more(self, bytes):
 
     def __gracefully_terminate(self):
         self.client_terminated = self.server_terminated = True
-        
+
         try:
             if not self.stream.closing:
                 self.closed(1006)
@@ -93,7 +102,7 @@ def data_provider():
                     yield "#" * i
 
             self.send(data_provider())
-            
+
             for i in range(0, 200, 25):
                 self.send("*" * i)
 
@@ -104,9 +113,9 @@ def received_message(self, m):
 
         def closed(self, code, reason=None):
             ioloop.IOLoop.instance().stop()
-                
+
     ws = MyClient('http://localhost:9000/ws', protocols=['http-only', 'chat'])
     ws.connect()
-        
+
     ioloop.IOLoop.instance().start()
 
