Merge branch 'master' of github.com:Lawouach/WebSocket-for-Python

Lawouach · Lawouach · commit 2d4312daaeba · 2012-03-22T18:54:30.000+01:00
diff --git a/ws4py/client/tornadoclient.py b/ws4py/client/tornadoclient.py
@@ -1,34 +1,48 @@
 # -*- 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)
         host, port = parts.netloc, 80
         if ':' in host:
             host, port = parts.netloc.split(':')
+        self.io.set_close_callback(self.__connection_refused)
         self.io.connect((host, int(port)), self.__send_handshake)
 
+    def __connection_refused(self, *args, **kwargs):
+        self.server_terminated = True
+        self.closed(1005, 'Connection refused')
+
     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)
 
@@ -41,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)
@@ -59,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)
@@ -88,7 +102,7 @@ def data_provider():
                     yield "#" * i
 
             self.send(data_provider())
-            
+
             for i in range(0, 200, 25):
                 self.send("*" * i)
 
@@ -99,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()
 
