def test_simple(self):
client = self._setup_connection()
frame = websockets.Frame.from_file(client.rfile)
assert frame.payload == b'server-foobar'
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'client-foobar')))
client.wfile.flush()
frame = websockets.Frame.from_file(client.rfile)
assert frame.payload == b'client-foobar'
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.BINARY, payload=b'\xde\xad\xbe\xef')))
client.wfile.flush()
frame = websockets.Frame.from_file(client.rfile)
assert frame.payload == b'\xde\xad\xbe\xef'
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE)))
client.wfile.flush()
assert len(self.master.state.flows) == 2
assert isinstance(self.master.state.flows[0], HTTPFlow)
assert isinstance(self.master.state.flows[1], WebSocketFlow)
assert len(self.master.state.flows[1].messages) == 5
assert self.master.state.flows[1].messages[0].content == b'server-foobar'
assert self.master.state.flows[1].messages[0].type == websockets.OPCODE.TEXT
assert self.master.state.flows[1].messages[1].content == b'client-foobar'
assert self.master.state.flows[1].messages[1].type == websockets.OPCODE.TEXT
assert self.master.state.flows[1].messages[2].content == b'client-foobar'
assert self.master.state.flows[1].messages[2].type == websockets.OPCODE.TEXT
assert self.master.state.flows[1].messages[3].content == b'\xde\xad\xbe\xef'
assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY
assert self.master.state.flows[1].messages[4].content == b'\xde\xad\xbe\xef'
assert self.master.state.flows[1].messages[4].type == websockets.OPCODE.BINARY
def test_change_payload(self):
class Addon:
def websocket_message(self, f):
f.messages[-1].content = "foo"
self.proxy.set_addons(Addon())
self.setup_connection()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'foo'
self.client.wfile.write(
bytes(
websockets.Frame(fin=1,
mask=1,
opcode=websockets.OPCODE.TEXT,
payload=b'self.client-foobar')))
self.client.wfile.flush()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'foo'
self.client.wfile.write(
bytes(
websockets.Frame(fin=1,
mask=1,
opcode=websockets.OPCODE.BINARY,
payload=b'\xde\xad\xbe\xef')))
self.client.wfile.flush()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'foo'
def test_pong(self):
self.setup_connection()
self.client.wfile.write(
bytes(
websockets.Frame(fin=1,
mask=1,
opcode=websockets.OPCODE.PING,
payload=b'foobar')))
self.client.wfile.flush()
frame = websockets.Frame.from_file(self.client.rfile)
websockets.Frame.from_file(self.client.rfile)
self.client.wfile.write(
bytes(
websockets.Frame(fin=1, mask=1,
opcode=websockets.OPCODE.CLOSE)))
self.client.wfile.flush()
assert frame.header.opcode == websockets.OPCODE.PONG
assert frame.payload == b'foobar'
for i in range(20):
if self.master.has_log("Pong Received from server", "info"):
break
time.sleep(0.01)
else:
raise AssertionError("No pong seen")
def handle_websockets(cls, rfile, wfile):
wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'server-foobar')))
wfile.flush()
frame = websockets.Frame.from_file(rfile)
wfile.write(bytes(websockets.Frame(fin=1, opcode=frame.header.opcode, payload=frame.payload)))
wfile.flush()
def handle_websockets(cls, rfile, wfile):
frame = websockets.Frame.from_file(rfile)
wfile.write(bytes(websockets.Frame(fin=1, opcode=frame.header.opcode, payload=frame.payload)))
wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE)))
wfile.flush()
with pytest.raises(exceptions.TcpDisconnect):
websockets.Frame.from_file(rfile)
def _handle_data_frame(self, frame, source_conn, other_conn, is_server):
fb = self.server_frame_buffer if is_server else self.client_frame_buffer
fb.append(frame)
if frame.header.fin:
payload = b''.join(f.payload for f in fb)
original_chunk_sizes = [len(f.payload) for f in fb]
message_type = fb[0].header.opcode
compressed_message = fb[0].header.rsv1
fb.clear()
websocket_message = WebSocketMessage(message_type, not is_server,
payload)
length = len(websocket_message.content)
self.flow.messages.append(websocket_message)
self.channel.ask("websocket_message", self.flow)
def get_chunk(payload):
if len(payload) == length:
# message has the same length, we can reuse the same sizes
pos = 0
for s in original_chunk_sizes:
yield payload[pos:pos + s]
pos += s
else:
# just re-chunk everything into 10kB frames
chunk_size = 10240
chunks = range(0, len(payload), chunk_size)
for i in chunks:
yield payload[i:i + chunk_size]
frms = [
websockets.Frame(
payload=chunk,
opcode=frame.header.opcode,
mask=(False if is_server else 1),
masking_key=(b'' if is_server else os.urandom(4)))
for chunk in get_chunk(websocket_message.content)
]
if len(frms) > 0:
frms[-1].header.fin = True
else:
frms.append(
websockets.Frame(
fin=True,
opcode=websockets.OPCODE.CONTINUE,
mask=(False if is_server else 1),
masking_key=(b'' if is_server else os.urandom(4))))
frms[0].header.opcode = message_type
frms[0].header.rsv1 = compressed_message
for frm in frms:
other_conn.send(bytes(frm))
return True
def handle_websockets(cls, rfile, wfile):
wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PING, payload=b'foobar')))
wfile.flush()
frame = websockets.Frame.from_file(rfile)
assert frame.header.opcode == websockets.OPCODE.PONG
assert frame.payload == b'foobar'
wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'pong-received')))
wfile.flush()
def handle_websockets(cls, rfile, wfile):
frame = websockets.Frame.from_file(rfile)
assert frame.header.opcode == websockets.OPCODE.PING
assert frame.payload == b''
wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PONG, payload=frame.payload)))
wfile.flush()
wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE)))
wfile.flush()
websockets.Frame.from_file(rfile)
def test_simple_tls(self):
self.setup_connection()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'server-foobar'
self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.TEXT, payload=b'self.client-foobar')))
self.client.wfile.flush()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'self.client-foobar'
self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE)))
self.client.wfile.flush()
def test_pong(self):
self.setup_connection()
self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.PING, payload=b'foobar')))
self.client.wfile.flush()
frame = websockets.Frame.from_file(self.client.rfile)
websockets.Frame.from_file(self.client.rfile)
self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE)))
self.client.wfile.flush()
assert frame.header.opcode == websockets.OPCODE.PONG
assert frame.payload == b'foobar'
assert self.master.has_log("Pong Received from server", "info")
def handle_websockets(cls, rfile, wfile):
wfile.write(
bytes(
websockets.Frame(fin=1,
opcode=websockets.OPCODE.TEXT,
payload=b'server-foobar')))
wfile.flush()
def handle_websockets(cls, rfile, wfile):
frame = websockets.Frame.from_file(rfile)
assert frame.header.opcode == websockets.OPCODE.TEXT
success = frame.payload == b'This is an injected message!'
wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=str(success).encode())))
wfile.flush()
def _handle_data_frame(self, frame, source_conn, other_conn, is_server):
fb = self.server_frame_buffer if is_server else self.client_frame_buffer
fb.append(frame)
if frame.header.fin:
if frame.header.opcode == websockets.OPCODE.TEXT:
t = WebSocketTextMessage
else:
t = WebSocketBinaryMessage
payload = b''.join(f.payload for f in fb)
fb.clear()
websocket_message = t(self.flow, not is_server, payload)
self.flow.messages.append(websocket_message)
self.channel.ask("websocket_message", self.flow)
# chunk payload into multiple 10kB frames, and send them
payload = websocket_message.content
chunk_size = 10240 # 10kB
chunks = range(0, len(payload), chunk_size)
frms = [
websockets.Frame(
payload=payload[i:i + chunk_size],
opcode=frame.header.opcode,
mask=(False if is_server else 1),
masking_key=(b'' if is_server else os.urandom(4)))
for i in chunks
]
frms[-1].header.fin = 1
for frm in frms:
other_conn.send(bytes(frm))
return True
def test_close_payload_2(self):
self.setup_connection()
self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42foobar')))
self.client.wfile.flush()
websockets.Frame.from_file(self.client.rfile)
with pytest.raises(exceptions.TcpDisconnect):
websockets.Frame.from_file(self.client.rfile)
def test_close(self):
client = self._setup_connection()
client.wfile.write(
bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE)))
client.wfile.flush()
with pytest.raises(exceptions.TcpDisconnect):
websockets.Frame.from_file(client.rfile)
def test_serialization_bijection(self, masked, length):
frame = websockets.Frame(
os.urandom(length),
fin=True,
opcode=websockets.OPCODE.TEXT,
mask=int(masked),
masking_key=(os.urandom(4) if masked else None))
serialized = bytes(frame)
assert frame == websockets.Frame.from_bytes(serialized)
def test_pong(self):
client = self._setup_connection()
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PING, payload=b'foobar')))
client.wfile.flush()
frame = websockets.Frame.from_file(client.rfile)
assert frame.header.opcode == websockets.OPCODE.PONG
assert frame.payload == b'foobar'
async def test_ping(self):
self.setup_connection()
frame = websockets.Frame.from_file(self.client.rfile)
websockets.Frame.from_file(self.client.rfile)
self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE)))
self.client.wfile.flush()
assert frame.header.opcode == websockets.OPCODE.PING
assert frame.payload == b'' # We don't send payload to other end
assert await self.master.await_log("Pong Received from server", "info")
def test_simple(self, streaming):
class Stream:
def websocket_start(self, f):
f.stream = streaming
self.master.addons.add(Stream())
self.setup_connection()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'server-foobar'
self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.TEXT, payload=b'self.client-foobar')))
self.client.wfile.flush()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'self.client-foobar'
self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.BINARY, payload=b'\xde\xad\xbe\xef')))
self.client.wfile.flush()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'\xde\xad\xbe\xef'
self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE)))
self.client.wfile.flush()
assert len(self.master.state.flows) == 2
assert isinstance(self.master.state.flows[0], HTTPFlow)
assert isinstance(self.master.state.flows[1], WebSocketFlow)
assert len(self.master.state.flows[1].messages) == 5
assert self.master.state.flows[1].messages[0].content == 'server-foobar'
assert self.master.state.flows[1].messages[0].type == websockets.OPCODE.TEXT
assert self.master.state.flows[1].messages[1].content == 'self.client-foobar'
assert self.master.state.flows[1].messages[1].type == websockets.OPCODE.TEXT
assert self.master.state.flows[1].messages[2].content == 'self.client-foobar'
assert self.master.state.flows[1].messages[2].type == websockets.OPCODE.TEXT
assert self.master.state.flows[1].messages[3].content == b'\xde\xad\xbe\xef'
assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY
assert self.master.state.flows[1].messages[4].content == b'\xde\xad\xbe\xef'
assert self.master.state.flows[1].messages[4].type == websockets.OPCODE.BINARY
def test_ping(self):
self.setup_connection()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.header.opcode == websockets.OPCODE.PING
assert frame.payload == b'foobar'
self.client.wfile.write(
bytes(
websockets.Frame(fin=1,
opcode=websockets.OPCODE.PONG,
payload=frame.payload)))
self.client.wfile.flush()
frame = websockets.Frame.from_file(self.client.rfile)
assert frame.header.opcode == websockets.OPCODE.TEXT
assert frame.payload == b'pong-received'
def test_equality(self):
f = websockets.Frame(payload=b'1234')
f2 = websockets.Frame(payload=b'1234')
assert f == f2
assert f != b'1234'
def test_human_readable(self):
f = websockets.Frame()
assert repr(f)
f = websockets.Frame(b"foobar")
assert "foobar" in repr(f)
def handle_websockets(cls, rfile, wfile):
wfile.write(
bytes(websockets.Frame(fin=1, opcode=15, payload=b'foobar')))
wfile.flush()
def round(*args, **kwargs):
f = websockets.Frame(*args, **kwargs)
raw = bytes(f)
f2 = websockets.Frame.from_file(tutils.treader(raw))
assert f == f2
