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)
