├── README.md ├── LICENSE ├── test_websocket.py └── websocket.py /README.md: -------------------------------------------------------------------------------- 1 | Websocket client for protocol version 13 using the Tornado IO loop. 2 | 3 | 4 | 5 | 6 | ```python 7 | import websocket 8 | 9 | class HelloSocket(websocket.WebSocket): 10 | 11 | def on_open(self): 12 | self.write('hello, world') 13 | 14 | def on_message(self, data): 15 | print data 16 | 17 | def on_ping(self): 18 | print 'I was pinged' 19 | 20 | def on_pong(self): 21 | print 'I was ponged' 22 | 23 | def on_close(self): 24 | print 'Socket closed.' 25 | 26 | 27 | ws = HelloSocket('ws://echo.websocket.org') 28 | ws.connect() 29 | ``` 30 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | Copyright (c) 2012 Jeff Balogh 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy of 5 | this software and associated documentation files (the "Software"), to deal in 6 | the Software without restriction, including without limitation the rights to 7 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 8 | of the Software, and to permit persons to whom the Software is furnished to do 9 | so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in all 12 | copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 19 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 20 | SOFTWARE. 21 | -------------------------------------------------------------------------------- /test_websocket.py: -------------------------------------------------------------------------------- 1 | #!python 2 | # coding=utf-8 3 | from functools import partial 4 | 5 | from tornado import testing 6 | import tornado.websocket 7 | import tornado.web 8 | import tornado.ioloop 9 | import websocket 10 | 11 | 12 | class EchoWebSocketHandler(tornado.websocket.WebSocketHandler): 13 | def on_message(self, message): 14 | self.write_message(message) 15 | 16 | 17 | class WebSocketTest(testing.AsyncHTTPTestCase): 18 | """ 19 | Example of WebSocket usage as a client 20 | in AsyncHTTPTestCase-based unit tests. 21 | """ 22 | 23 | def get_app(self): 24 | app = tornado.web.Application([('/', EchoWebSocketHandler)]) 25 | return app 26 | 27 | def test_echo(self): 28 | _self = self 29 | 30 | class WSClient(websocket.WebSocket): 31 | def on_open(self): 32 | self.write_message('hello') 33 | 34 | def on_message(self, data): 35 | _self.assertEquals(data, 'hello') 36 | _self.io_loop.add_callback(_self.stop) 37 | 38 | self.io_loop.add_callback(partial(WSClient, self.get_url('/'), 39 | self.io_loop)) 40 | self.wait() 41 | -------------------------------------------------------------------------------- /websocket.py: -------------------------------------------------------------------------------- 1 | """ 2 | Websocket client for protocol version 13 using the Tornado IO loop. 3 | 4 | http://tools.ietf.org/html/rfc6455 5 | 6 | Based on the websocket server in tornado/websocket.py by Jacob Kristhammar. 7 | """ 8 | import array 9 | import base64 10 | import functools 11 | import hashlib 12 | import logging 13 | import os 14 | import re 15 | import socket 16 | import struct 17 | import sys 18 | import time 19 | import urlparse 20 | 21 | import tornado.escape 22 | from tornado import ioloop, iostream 23 | from tornado.httputil import HTTPHeaders 24 | from tornado.util import bytes_type, b 25 | 26 | 27 | # The initial handshake over HTTP. 28 | INIT = """\ 29 | GET %(path)s HTTP/1.1 30 | Host: %(host)s:%(port)s 31 | Upgrade: websocket 32 | Connection: Upgrade 33 | Sec-Websocket-Key: %(key)s 34 | Sec-Websocket-Version: 13\ 35 | """ 36 | 37 | # Magic string defined in the spec for calculating keys. 38 | MAGIC = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' 39 | 40 | 41 | def frame(data, opcode=0x01): 42 | """Encode data in a websocket frame.""" 43 | # [fin, rsv, rsv, rsv] [opcode] 44 | frame = struct.pack('B', 0x80 | opcode) 45 | 46 | # Our next bit is 1 since we're using a mask. 47 | length = len(data) 48 | if length < 126: 49 | # If length < 126, it fits in the next 7 bits. 50 | frame += struct.pack('B', 0x80 | length) 51 | elif length <= 0xFFFF: 52 | # If length < 0xffff, put 126 in the next 7 bits and write the length 53 | # in the next 2 bytes. 54 | frame += struct.pack('!BH', 0x80 | 126, length) 55 | else: 56 | # Otherwise put 127 in the next 7 bits and write the length in the next 57 | # 8 bytes. 58 | frame += struct.pack('!BQ', 0x80 | 127, length) 59 | 60 | # Clients must apply a 32-bit mask to all data sent. 61 | mask = map(ord, os.urandom(4)) 62 | frame += struct.pack('!BBBB', *mask) 63 | # Mask each byte of data using a byte from the mask. 64 | msg = [ord(c) ^ mask[i % 4] for i, c in enumerate(data)] 65 | frame += struct.pack('!' + 'B' * length, *msg) 66 | return frame 67 | 68 | 69 | class WebSocket(object): 70 | """Websocket client for protocol version 13 using the Tornado IO loop.""" 71 | 72 | def __init__(self, url, io_loop=None, extra_headers=None): 73 | ports = {'ws': 80, 'wss': 443} 74 | 75 | self.url = urlparse.urlparse(url) 76 | self.host = self.url.hostname 77 | self.port = self.url.port or ports[self.url.scheme] 78 | self.path = self.url.path or '/' 79 | 80 | if extra_headers is not None and len(extra_headers) > 0: 81 | header_set = [] 82 | for k, v in extra_headers.iteritems(): 83 | header_set.append("%s: %s" % (k, v)) 84 | self.headers = "\r\n".join(header_set) 85 | else: 86 | self.headers = None 87 | 88 | self.client_terminated = False 89 | self.server_terminated = False 90 | self._final_frame = False 91 | self._frame_opcode = None 92 | self._frame_length = None 93 | self._fragmented_message_buffer = None 94 | self._fragmented_message_opcode = None 95 | self._waiting = None 96 | 97 | self.key = base64.b64encode(os.urandom(16)) 98 | self.stream = iostream.IOStream(socket.socket(), io_loop) 99 | self.stream.connect((self.host, self.port), self._on_connect) 100 | 101 | def on_open(self): 102 | pass 103 | 104 | def on_message(self, data): 105 | pass 106 | 107 | def on_ping(self): 108 | pass 109 | 110 | def on_pong(self): 111 | pass 112 | 113 | def on_close(self): 114 | pass 115 | 116 | def on_unsupported(self): 117 | pass 118 | 119 | def write_message(self, message, binary=False): 120 | """Sends the given message to the client of this Web Socket.""" 121 | if binary: 122 | opcode = 0x2 123 | else: 124 | opcode = 0x1 125 | message = tornado.escape.utf8(message) 126 | assert isinstance(message, bytes_type) 127 | self._write_frame(True, opcode, message) 128 | 129 | def ping(self): 130 | self._write_frame(True, 0x9, b('')) 131 | 132 | def close(self): 133 | """Closes the WebSocket connection.""" 134 | if not self.server_terminated: 135 | if not self.stream.closed(): 136 | self._write_frame(True, 0x8, b("")) 137 | self.server_terminated = True 138 | if self.client_terminated: 139 | if self._waiting is not None: 140 | self.stream.io_loop.remove_timeout(self._waiting) 141 | self._waiting = None 142 | self.stream.close() 143 | elif self._waiting is None: 144 | # Give the client a few seconds to complete a clean shutdown, 145 | # otherwise just close the connection. 146 | self._waiting = self.stream.io_loop.add_timeout( 147 | time.time() + 5, self._abort) 148 | 149 | def _write_frame(self, fin, opcode, data): 150 | self.stream.write(frame(data, opcode)) 151 | 152 | def _on_connect(self): 153 | request = '\r\n'.join(INIT.splitlines()) % self.__dict__ 154 | if self.headers is not None: 155 | request += '\r\n' + self.headers 156 | request += '\r\n\r\n' 157 | self.stream.write(tornado.escape.utf8(request)) 158 | self.stream.read_until('\r\n\r\n', self._on_headers) 159 | 160 | def _on_headers(self, data): 161 | first, _, rest = data.partition('\r\n') 162 | headers = HTTPHeaders.parse(rest) 163 | # Expect HTTP 101 response. 164 | if not re.match('HTTP/[^ ]+ 101', first): 165 | self._async_callback(self.on_unsupported)() 166 | self.close() 167 | else: 168 | # Expect Connection: Upgrade. 169 | assert headers['Connection'].lower() == 'upgrade' 170 | # Expect Upgrade: websocket. 171 | assert headers['Upgrade'].lower() == 'websocket' 172 | # Sec-WebSocket-Accept should be derived from our key. 173 | accept = base64.b64encode(hashlib.sha1(self.key + MAGIC).digest()) 174 | assert headers['Sec-WebSocket-Accept'] == accept 175 | 176 | self._async_callback(self.on_open)() 177 | self._receive_frame() 178 | 179 | def _receive_frame(self): 180 | self.stream.read_bytes(2, self._on_frame_start) 181 | 182 | def _on_frame_start(self, data): 183 | header, payloadlen = struct.unpack("BB", data) 184 | self._final_frame = header & 0x80 185 | reserved_bits = header & 0x70 186 | self._frame_opcode = header & 0xf 187 | self._frame_opcode_is_control = self._frame_opcode & 0x8 188 | if reserved_bits: 189 | # client is using as-yet-undefined extensions; abort 190 | return self._abort() 191 | if (payloadlen & 0x80): 192 | # Masked frame -> abort connection 193 | return self._abort() 194 | payloadlen = payloadlen & 0x7f 195 | if self._frame_opcode_is_control and payloadlen >= 126: 196 | # control frames must have payload < 126 197 | return self._abort() 198 | if payloadlen < 126: 199 | self._frame_length = payloadlen 200 | self.stream.read_bytes(self._frame_length, self._on_frame_data) 201 | elif payloadlen == 126: 202 | self.stream.read_bytes(2, self._on_frame_length_16) 203 | elif payloadlen == 127: 204 | self.stream.read_bytes(8, self._on_frame_length_64) 205 | 206 | def _on_frame_length_16(self, data): 207 | self._frame_length = struct.unpack("!H", data)[0] 208 | self.stream.read_bytes(self._frame_length, self._on_frame_data) 209 | 210 | def _on_frame_length_64(self, data): 211 | self._frame_length = struct.unpack("!Q", data)[0] 212 | self.stream.read_bytes(self._frame_length, self._on_frame_data) 213 | 214 | def _on_frame_data(self, data): 215 | unmasked = array.array("B", data) 216 | 217 | if self._frame_opcode_is_control: 218 | # control frames may be interleaved with a series of fragmented 219 | # data frames, so control frames must not interact with 220 | # self._fragmented_* 221 | if not self._final_frame: 222 | # control frames must not be fragmented 223 | self._abort() 224 | return 225 | opcode = self._frame_opcode 226 | elif self._frame_opcode == 0: # continuation frame 227 | if self._fragmented_message_buffer is None: 228 | # nothing to continue 229 | self._abort() 230 | return 231 | self._fragmented_message_buffer += unmasked 232 | if self._final_frame: 233 | opcode = self._fragmented_message_opcode 234 | unmasked = self._fragmented_message_buffer 235 | self._fragmented_message_buffer = None 236 | else: # start of new data message 237 | if self._fragmented_message_buffer is not None: 238 | # can't start new message until the old one is finished 239 | self._abort() 240 | return 241 | if self._final_frame: 242 | opcode = self._frame_opcode 243 | else: 244 | self._fragmented_message_opcode = self._frame_opcode 245 | self._fragmented_message_buffer = unmasked 246 | 247 | if self._final_frame: 248 | self._handle_message(opcode, unmasked.tostring()) 249 | 250 | if not self.client_terminated: 251 | self._receive_frame() 252 | 253 | def _abort(self): 254 | """Instantly aborts the WebSocket connection by closing the socket""" 255 | self.client_terminated = True 256 | self.server_terminated = True 257 | self.stream.close() 258 | self.close() 259 | 260 | def _handle_message(self, opcode, data): 261 | if self.client_terminated: 262 | return 263 | 264 | if opcode == 0x1: 265 | # UTF-8 data 266 | try: 267 | decoded = data.decode("utf-8") 268 | except UnicodeDecodeError: 269 | self._abort() 270 | return 271 | self._async_callback(self.on_message)(decoded) 272 | elif opcode == 0x2: 273 | # Binary data 274 | self._async_callback(self.on_message)(data) 275 | elif opcode == 0x8: 276 | # Close 277 | self.client_terminated = True 278 | self.close() 279 | elif opcode == 0x9: 280 | # Ping 281 | self._write_frame(True, 0xA, data) 282 | self._async_callback(self.on_ping)() 283 | elif opcode == 0xA: 284 | # Pong 285 | self._async_callback(self.on_pong)() 286 | else: 287 | self._abort() 288 | 289 | def _async_callback(self, callback, *args, **kwargs): 290 | """Wrap callbacks with this if they are used on asynchronous requests. 291 | 292 | Catches exceptions properly and closes this WebSocket if an exception 293 | is uncaught. 294 | """ 295 | if args or kwargs: 296 | callback = functools.partial(callback, *args, **kwargs) 297 | 298 | def wrapper(*args, **kwargs): 299 | try: 300 | return callback(*args, **kwargs) 301 | except Exception: 302 | logging.error('Uncaught exception', exc_info=True) 303 | self._abort() 304 | return wrapper 305 | 306 | 307 | def main(url, message='hello, world'): 308 | 309 | class HelloSocket(WebSocket): 310 | 311 | def on_open(self): 312 | self.ping() 313 | print '>>', message 314 | self.write_message(message) 315 | 316 | def on_message(self, data): 317 | print 'on_message:', data 318 | msg = raw_input('>> ') 319 | if msg == 'ping': 320 | self.ping() 321 | elif msg == 'die': 322 | self.close() 323 | else: 324 | self.write_message(msg) 325 | 326 | def on_close(self): 327 | print 'on_close' 328 | 329 | def on_pong(self): 330 | print 'on_pong' 331 | 332 | ws = HelloSocket(url) 333 | try: 334 | ioloop.IOLoop.instance().start() 335 | except KeyboardInterrupt: 336 | pass 337 | finally: 338 | ws.close() 339 | 340 | 341 | if __name__ == '__main__': 342 | main(*sys.argv[1:]) 343 | --------------------------------------------------------------------------------