def handle_websockets(cls, rfile, wfile):
wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.TEXT, payload=b'server-foobar')))
wfile.flush()
header, frame, _ = websocket_utils.read_raw_frame(rfile)
wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=header.opcode, payload=frame.payload)))
wfile.flush()
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, _ = websocket.read_frame(self.client.rfile)
assert frame.payload == b'foo'
self.client.wfile.write(
bytes(
websockets_frame.Frame(fin=1,
mask=1,
opcode=Opcode.TEXT,
payload=b'self.client-foobar')))
self.client.wfile.flush()
_, frame, _ = websocket.read_frame(self.client.rfile)
assert frame.payload == b'foo'
self.client.wfile.write(
bytes(
websockets_frame.Frame(fin=1,
mask=1,
opcode=Opcode.BINARY,
payload=b'\xde\xad\xbe\xef')))
self.client.wfile.flush()
_, frame, _ = websocket.read_frame(self.client.rfile)
assert frame.payload == b'foo'
def handle_websockets(cls, rfile, wfile):
header, frame, _ = websocket_utils.read_raw_frame(rfile)
wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=header.opcode, payload=frame.payload)))
wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.CLOSE)))
wfile.flush()
with pytest.raises(exceptions.TcpDisconnect):
_, _, _ = websocket_utils.read_raw_frame(rfile)
def test_simple(self, streaming):
class Stream:
def websocket_start(self, f):
f.stream = streaming
self.proxy.set_addons(Stream())
self.setup_connection()
_, frame, _ = websocket.read_frame(self.client.rfile)
assert frame.payload == b'server-foobar'
self.client.wfile.write(
bytes(
websockets_frame.Frame(fin=1,
mask=1,
opcode=Opcode.TEXT,
payload=b'self.client-foobar')))
self.client.wfile.flush()
_, frame, _ = websocket.read_frame(self.client.rfile)
assert frame.payload == b'self.client-foobar'
self.client.wfile.write(
bytes(
websockets_frame.Frame(fin=1,
mask=1,
opcode=Opcode.BINARY,
payload=b'\xde\xad\xbe\xef')))
self.client.wfile.flush()
_, frame, _ = websocket.read_frame(self.client.rfile)
assert frame.payload == b'\xde\xad\xbe\xef'
self.client.wfile.write(
bytes(websockets_frame.Frame(fin=1, mask=1, opcode=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 == Opcode.TEXT
assert self.master.state.flows[1].messages[
1].content == 'self.client-foobar'
assert self.master.state.flows[1].messages[1].type == Opcode.TEXT
assert self.master.state.flows[1].messages[
2].content == 'self.client-foobar'
assert self.master.state.flows[1].messages[2].type == 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 == 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 == Opcode.BINARY
def handle_websockets(cls, rfile, wfile):
header, frame, _ = websocket_utils.read_raw_frame(rfile)
assert header.opcode == Opcode.PING
assert frame.payload == b''
wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.PONG, payload=frame.payload)))
wfile.flush()
wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.CLOSE)))
wfile.flush()
_ = websocket_utils.read_raw_frame(rfile)
def test_simple_tls(self):
self.setup_connection()
_, frame, _ = websocket_utils.read_raw_frame(self.client.rfile)
assert frame.payload == b'server-foobar'
self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.TEXT, payload=b'self.client-foobar')))
self.client.wfile.flush()
_, frame, _ = websocket_utils.read_raw_frame(self.client.rfile)
assert frame.payload == b'self.client-foobar'
self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.CLOSE)))
self.client.wfile.flush()
async def test_pong(self):
self.setup_connection()
self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.PING, payload=b'foobar')))
self.client.wfile.flush()
header, frame, _ = websocket_utils.read_raw_frame(self.client.rfile)
_ = websocket_utils.read_raw_frame(self.client.rfile)
self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.CLOSE)))
self.client.wfile.flush()
assert header.opcode == Opcode.PONG
assert frame.payload == b'foobar'
assert await self.master.await_log("pong received")
def handle_websockets(cls, rfile, wfile):
wfile.write(
bytes(
websockets_frame.Frame(fin=1,
opcode=Opcode.TEXT,
payload=b'server-foobar')))
wfile.flush()
def handle_websockets(cls, rfile, wfile):
header, frame, _ = websocket_utils.read_raw_frame(rfile)
assert header.opcode == Opcode.TEXT
success = frame.payload == b'This is an injected message!'
wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.TEXT, payload=str(success).encode())))
wfile.flush()
def test_serialization_bijection(self, masked, length):
frame = websockets_frame.Frame(
os.urandom(length),
fin=True,
opcode=Opcode.TEXT,
mask=int(masked),
masking_key=(os.urandom(4) if masked else None))
serialized = bytes(frame)
assert frame == websockets_frame.Frame.from_bytes(serialized)
def test_close_payload_2(self):
self.setup_connection()
self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.CLOSE, payload=b'\00\42foobar')))
self.client.wfile.flush()
_ = websocket_utils.read_raw_frame(self.client.rfile)
with pytest.raises(exceptions.TcpDisconnect):
_ = websocket_utils.read_raw_frame(self.client.rfile)
def test_close(self):
self.setup_connection()
self.client.wfile.write(
bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.CLOSE)))
self.client.wfile.flush()
_ = websocket.read_frame(self.client.rfile)
with pytest.raises(exceptions.TcpDisconnect):
_ = websocket.read_frame(self.client.rfile)
async def test_ping(self):
self.setup_connection()
header, frame, _ = websocket_utils.read_raw_frame(self.client.rfile)
_ = websocket_utils.read_raw_frame(self.client.rfile)
self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.CLOSE)))
self.client.wfile.flush()
assert header.opcode == 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 handle_websockets(cls, rfile, wfile):
wfile.write(
bytes(
websockets_frame.Frame(fin=1,
opcode=Opcode.PING,
payload=b'foobar')))
wfile.flush()
header, frame, _ = websocket.read_frame(rfile)
assert header.opcode == Opcode.PONG
assert frame.payload == b'foobar'
wfile.write(
bytes(
websockets_frame.Frame(fin=1,
opcode=Opcode.PONG,
payload=b'done')))
wfile.flush()
wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.CLOSE)))
wfile.flush()
_ = websocket.read_frame(rfile, False)
def test_equality(self):
f = websockets_frame.Frame(payload=b'1234')
f2 = websockets_frame.Frame(payload=b'1234')
assert f == f2
assert f != b'1234'
def test_human_readable(self):
f = websockets_frame.Frame()
assert repr(f)
f = websockets_frame.Frame(b"foobar")
assert "foobar" in repr(f)
def round(*args, **kwargs):
f = websockets_frame.Frame(*args, **kwargs)
raw = bytes(f)
f2 = websockets_frame.Frame.from_file(tutils.treader(raw))
assert f == f2
