1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81
| import struct import SocketServer from base64 import b64encode, b64decode from hashlib import sha1 from mimetools import Message from StringIO import StringIO clients = []; class WebSocketsHandler(SocketServer.StreamRequestHandler): magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' def setup(self): SocketServer.StreamRequestHandler.setup(self) print "connection established", self.client_address self.handshake_done = False def handle(self): while True: if not self.handshake_done: self.handshake() clients.append(self) else: try: self.read_next_message() except: self.request.close() return def read_next_message(self): length = ord(self.rfile.read(2)[1]) & 127 if length == 126: length = struct.unpack(">H", self.rfile.read(2))[0] elif length == 127: length = struct.unpack(">Q", self.rfile.read(8))[0] masks = [ord(byte) for byte in self.rfile.read(4)] decoded = "" for char in self.rfile.read(length): decoded += chr(ord(char) ^ masks[len(decoded) % 4]) self.on_message(decoded) def send_message(self, data, fin=True, opcode=1, masking_key=False): header = struct.pack('!B', ((fin << 7) | (0 << 6)| (0 << 5)| (0 << 4)| opcode)) if masking_key: mask_bit = 1 << 7 else: mask_bit = 0 length = len(data) if length < 126: header += struct.pack('!B', (mask_bit | length)) elif length < (1 << 16): header += struct.pack('!B', (mask_bit | 126)) + struct.pack('!H', length) elif length < (1 << 63): header += struct.pack('!B', (mask_bit | 127)) + struct.pack('!Q', length) body = data self.request.send(bytes(header + body)) def handshake(self): data = self.request.recv(1024).strip() headers = Message(StringIO(data.split('\r\n', 1)[1])) if headers.get("Upgrade", None) != "websocket": return print 'Handshaking...' key = headers['Sec-WebSocket-Key'] digest = b64encode(sha1(key + self.magic).hexdigest().decode('hex')) response = 'HTTP/1.1 101 Switching Protocols\r\n'+'Upgrade: websocket\r\n'+'Connection: Upgrade\r\n' response += 'Sec-WebSocket-Accept: %s\r\n\r\n' % digest self.handshake_done = self.request.send(response) def on_message(self, message): text = b64decode(message) print text for client in clients: if client.client_address != self.client_address: client.send_message(b64encode(text)) class ThreadedTCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer): pass if __name__ == "__main__": server = ThreadedTCPServer(("localhost", 9999), WebSocketsHandler) server.serve_forever()
|