def test_send_text_masked(self): writer = websocket.WebSocketWriter(self.transport, use_mask=True, random=random.Random(123)) writer.send(b'text') self.transport.write.assert_called_with( b'\x81\x84\rg\xb3fy\x02\xcb\x12')
def heartbeat(self): # setup pipes read_transport, read_proto = yield from self.loop.connect_read_pipe( aiohttp.StreamProtocol, os.fdopen(self.up_read, 'rb')) write_transport, _ = yield from self.loop.connect_write_pipe( aiohttp.StreamProtocol, os.fdopen(self.down_write, 'wb')) reader = read_proto.reader.set_parser(websocket.WebSocketParser) writer = websocket.WebSocketWriter(write_transport) while True: try: msg = yield from reader.read() except: print('Supervisor is dead, {} stopping...'.format(os.getpid())) self.loop.stop() break if msg.tp == websocket.MSG_PING: writer.pong() elif msg.tp == websocket.MSG_CLOSE: break read_transport.close() write_transport.close()
def heartbeat(self): # setup pipes read_transport, read_proto = yield from self.loop.connect_read_pipe( aiohttp.StreamProtocol, os.fdopen(self.up_read, 'rb')) write_transport, _ = yield from self.loop.connect_write_pipe( aiohttp.StreamProtocol, os.fdopen(self.down_write, 'wb')) reader = read_proto.reader.set_parser(websocket.WebSocketParser) writer = websocket.WebSocketWriter(write_transport) asyncio.Task(self.start_server(writer)) while True: try: msg = yield from reader.read() except: print('Supervisor is dead, {} stopping...'.format(os.getpid())) self.loop.stop() break if msg.tp == websocket.MSG_PING: writer.pong() elif msg.tp == websocket.MSG_CLOSE: break elif msg.tp == websocket.MSG_TEXT: # broadcast message for wsc in self.clients: wsc.send(msg.data.strip().encode()) read_transport.close() write_transport.close()
def connect_ws(self, url, protocol='chat'): sec_key = base64.b64encode(os.urandom(16)) conn = aiohttp.TCPConnector(loop=self.loop) self.addCleanup(conn.close) # send request response = yield from aiohttp.request( 'get', url, headers={ 'UPGRADE': 'WebSocket', 'CONNECTION': 'Upgrade', 'SEC-WEBSOCKET-VERSION': '13', 'SEC-WEBSOCKET-PROTOCOL': protocol, 'SEC-WEBSOCKET-KEY': sec_key.decode(), }, connector=conn, loop=self.loop) self.addCleanup(response.close, True) self.assertEqual(101, response.status) self.assertEqual(response.headers.get('upgrade', '').lower(), 'websocket') self.assertEqual(response.headers.get('connection', '').lower(), 'upgrade') key = response.headers.get('sec-websocket-accept', '').encode() match = base64.b64encode(hashlib.sha1(sec_key + WS_KEY).digest()) self.assertEqual(key, match) # switch to websocket protocol connection = response.connection reader = connection.reader.set_parser(websocket.WebSocketParser) writer = websocket.WebSocketWriter(connection.writer) return response, reader, writer
def _handshake(self, response, sec_key): """Websocket handshake""" # Using WS_KEY in handshake key = response.headers.get('sec-websocket-accept', '').encode() match = base64.b64encode(hashlib.sha1(sec_key + WS_KEY).digest()) if key != match: raise ValueError("Handshake error - Invalid challenge response") # switch to websocket protocol self.connection = response.connection self.stream = self.connection.reader.set_parser( websocket.WebSocketParser) self.writer = websocket.WebSocketWriter(self.connection.writer) self.response = response
def connect(self, pid, up_write, down_read): read_transport, proto = yield from self.loop.connect_read_pipe( aiohttp.StreamProtocol, os.fdopen(down_read, 'rb')) write_transport, _ = yield from self.loop.connect_write_pipe( aiohttp.StreamProtocol, os.fdopen(up_write, 'wb')) reader = proto.reader.set_parser(websocket.WebSocketParser) writer = websocket.WebSocketWriter(write_transport) # store info self.pid = pid self.ping = time.monotonic() self.writer = writer self.rtransport = read_transport self.wtransport = write_transport self.chat_task = asyncio. async (self.chat(reader)) self.heartbeat_task = asyncio. async (self.heartbeat(writer))
def start_client(loop, url): name = input('Please enter your name: ').encode() sec_key = base64.b64encode(os.urandom(16)) # send request response = yield from aiohttp.request('get', url, headers={ 'UPGRADE': 'WebSocket', 'CONNECTION': 'Upgrade', 'SEC-WEBSOCKET-VERSION': '13', 'SEC-WEBSOCKET-KEY': sec_key.decode(), }, timeout=1.0) # websocket handshake if response.status != 101: raise ValueError("Handshake error: Invalid response status") if response.get('upgrade', '').lower() != 'websocket': raise ValueError("Handshake error - Invalid upgrade header") if response.get('connection', '').lower() != 'upgrade': raise ValueError("Handshake error - Invalid connection header") key = response.get('sec-websocket-accept', '').encode() match = base64.b64encode(hashlib.sha1(sec_key + WS_KEY).digest()) if key != match: raise ValueError("Handshake error - Invalid challenge response") # switch to websocket protocol stream = response.stream.set_parser(websocket.WebSocketParser) writer = websocket.WebSocketWriter(response.transport) # input reader def stdin_callback(): line = sys.stdin.buffer.readline() if not line: loop.stop() else: writer.send(name + b': ' + line) loop.add_reader(sys.stdin.fileno(), stdin_callback) @asyncio.coroutine def dispatch(): while True: try: msg = yield from stream.read() except aiohttp.EofStream: # server disconnected break if msg.tp == websocket.MSG_PING: writer.pong() elif msg.tp == websocket.MSG_TEXT: print(msg.data.strip()) elif msg.tp == websocket.MSG_CLOSE: break yield from dispatch()
def setUp(self): self.transport = unittest.mock.Mock() self.writer = websocket.WebSocketWriter(self.transport)
def connect(self, url, app_id, app_key, use_plaintext=True): date = datetime.datetime.utcnow() sec_key = base64.b64encode(os.urandom(16)) if use_plaintext: params = {'app_id': app_id, 'algorithm': 'key', 'app_key': binascii.hexlify(app_key)} else: datestr = date.replace(microsecond=0).isoformat() params = { 'date': datestr, 'app_id': app_id, 'algorithm': 'HMAC-SHA-256', 'signature': hmac.new(app_key,datestr.encode('ascii')+b' '+app_id.encode('utf-8'),hashlib.sha256).hexdigest() } response = yield from aiohttp.request( 'get', url + '?' + urllib.parse.urlencode(params), headers={ 'UPGRADE': 'WebSocket', 'CONNECTION': 'Upgrade', 'SEC-WEBSOCKET-VERSION': '13', 'SEC-WEBSOCKET-KEY': sec_key.decode(), }) if response.status == 401 and not use_plaintext: if 'Date' in response.headers: server_date = email.utils.parsedate_to_datetime(response.headers['Date']) if server_date.tzinfo is not None: server_date = (server_date - server_date.utcoffset()).replace(tzinfo=None) else: server_date = yield from response.read() server_date = datetime.datetime.strptime(server_date[:19].decode('ascii'), "%Y-%m-%dT%H:%M:%S") # Use delta on future requests date_delta = server_date - date print("Retrying authorization (delta=%s)" % date_delta) datestr = (date+date_delta).replace(microsecond=0).isoformat() params = { 'date': datestr, 'algorithm': 'HMAC-SHA-256', 'app_id': app_id, 'signature': hmac.new(app_key,datestr.encode('ascii')+b' '+app_id.encode('utf-8'),hashlib.sha256).hexdigest() } response = yield from aiohttp.request( 'get', url + '?' + urllib.parse.urlencode(params), headers={ 'UPGRADE': 'WebSocket', 'CONNECTION': 'Upgrade', 'SEC-WEBSOCKET-VERSION': '13', 'SEC-WEBSOCKET-KEY': sec_key.decode(), }) if response.status != 101: info = "%s %s\n" % (response.status, response.reason) for (k,v) in response.headers.items(): info += '%s: %s\n' % (k,v) info += '\n%s' % (yield from response.read()).decode('utf-8') if response.status == 401: raise RuntimeError("Authorization failure:\n%s" % info) elif response.status >= 500 and response.status < 600: raise RuntimeError("Server error:\n%s" % info) elif response.headers.get('upgrade', '').lower() != 'websocket': raise ValueError("Handshake error - Invalid upgrade header") elif response.headers.get('connection', '').lower() != 'upgrade': raise ValueError("Handshake error - Invalid connection header") else: raise ValueError("Handshake error: Invalid response status:\n%s" % info) key = response.headers.get('sec-websocket-accept', '').encode() match = base64.b64encode(hashlib.sha1(sec_key + WS_KEY).digest()) if key != match: raise ValueError("Handshake error - Invalid challenge response") # Switch to websocket protocol self.connection = response.connection self.stream = self.connection.reader.set_parser(websocket.WebSocketParser) self.writer = websocket.WebSocketWriter(self.connection.writer) self.response = response
def writer(transport): return websocket.WebSocketWriter(transport, use_mask=False)