def test_inject_message(ws_testdata):
tctx, playbook, flow = ws_testdata
assert (playbook << websocket.WebsocketStartHook(flow) >> reply() >>
WebSocketMessageInjected(
flow, WebSocketMessage(Opcode.TEXT, False, b"hello")) <<
websocket.WebsocketMessageHook(flow))
assert flow.websocket.messages[-1].content == b"hello"
assert flow.websocket.messages[-1].from_client is False
assert (playbook >> reply() << SendData(tctx.client, b"\x81\x05hello"))
def inject_websocket(self, flow: Flow, to_client: bool, message: bytes, is_text: bool = True):
if not isinstance(flow, http.HTTPFlow) or not flow.websocket:
ctx.log.warn("Cannot inject WebSocket messages into non-WebSocket flows.")
msg = websocket.WebSocketMessage(
Opcode.TEXT if is_text else Opcode.BINARY,
not to_client,
message
)
event = WebSocketMessageInjected(flow, msg)
try:
self.inject_event(event)
except ValueError as e:
ctx.log.warn(str(e))