def test_open_connection(tctx):
"""
If there is no server connection yet, establish one,
because the server may send data first.
"""
assert (Playbook(tcp.TCPLayer(tctx, True)) << OpenConnection(tctx.server))
tctx.server.state = ConnectionState.OPEN
assert (Playbook(tcp.TCPLayer(tctx, True)) << None)
def test_open_connection(tctx):
"""
If there is no server connection yet, establish one,
because the server may send data first.
"""
assert (
Playbook(tcp.TCPLayer(tctx, True))
<< OpenConnection(tctx.server)
)
tctx.server.timestamp_start = 1624544785
assert (
Playbook(tcp.TCPLayer(tctx, True))
<< None
)
def send_response(self, already_streamed: bool = False):
"""We have either consumed the entire response from the server or the response was set by an addon."""
assert self.flow.response
self.flow.response.timestamp_end = time.time()
yield HttpResponseHook(self.flow)
self.server_state = self.state_done
if (yield from self.check_killed(False)):
return
if not already_streamed:
content = self.flow.response.raw_content
yield SendHttp(ResponseHeaders(self.stream_id, self.flow.response, not content), self.context.client)
if content:
yield SendHttp(ResponseData(self.stream_id, content), self.context.client)
yield SendHttp(ResponseEndOfMessage(self.stream_id), self.context.client)
if self.flow.response.status_code == 101:
is_websocket = (
self.flow.response.headers.get("upgrade", "").lower() == "websocket"
and
self.flow.request.headers.get("Sec-WebSocket-Version", "") == "13"
)
if is_websocket:
self.child_layer = websocket.WebsocketLayer(self.context, self.flow)
else:
self.child_layer = tcp.TCPLayer(self.context)
if self.debug:
yield commands.Log(f"{self.debug}[http] upgrading to {self.child_layer}", "debug")
yield from self.child_layer.handle_event(events.Start())
self._handle_event = self.passthrough
return
def test_simple(tctx):
"""open connection, receive data, send it to peer"""
f = Placeholder(TCPFlow)
assert (
Playbook(tcp.TCPLayer(tctx))
<< tcp.TcpStartHook(f)
>> reply()
<< OpenConnection(tctx.server)
>> reply(None)
>> DataReceived(tctx.client, b"hello!")
<< tcp.TcpMessageHook(f)
>> reply()
<< SendData(tctx.server, b"hello!")
>> DataReceived(tctx.server, b"hi")
<< tcp.TcpMessageHook(f)
>> reply()
<< SendData(tctx.client, b"hi")
>> ConnectionClosed(tctx.server)
<< CloseConnection(tctx.client, half_close=True)
>> ConnectionClosed(tctx.client)
<< CloseConnection(tctx.server)
<< tcp.TcpEndHook(f)
>> reply()
>> ConnectionClosed(tctx.client)
<< None
)
assert len(f().messages) == 2
def send_response(self, already_streamed: bool = False):
"""We have either consumed the entire response from the server or the response was set by an addon."""
assert self.flow.response
self.flow.response.timestamp_end = time.time()
is_websocket = (self.flow.response.status_code == 101
and self.flow.response.headers.get(
"upgrade", "").lower() == "websocket"
and self.flow.request.headers.get(
"Sec-WebSocket-Version",
"").encode() == wsproto.handshake.WEBSOCKET_VERSION
and self.context.options.websocket)
if is_websocket:
# We need to set this before calling the response hook
# so that addons can determine if a WebSocket connection is following up.
self.flow.websocket = WebSocketData()
yield HttpResponseHook(self.flow)
self.server_state = self.state_done
if (yield from self.check_killed(False)):
return
if not already_streamed:
content = self.flow.response.raw_content
done_after_headers = not (content or self.flow.response.trailers)
yield SendHttp(
ResponseHeaders(self.stream_id, self.flow.response,
done_after_headers), self.context.client)
if content:
yield SendHttp(ResponseData(self.stream_id, content),
self.context.client)
if self.flow.response.trailers:
yield SendHttp(
ResponseTrailers(self.stream_id, self.flow.response.trailers),
self.context.client)
yield SendHttp(ResponseEndOfMessage(self.stream_id),
self.context.client)
if self.flow.response.status_code == 101:
if is_websocket:
self.child_layer = websocket.WebsocketLayer(
self.context, self.flow)
elif self.context.options.rawtcp:
self.child_layer = tcp.TCPLayer(self.context)
else:
yield commands.Log(
f"Sent HTTP 101 response, but no protocol is enabled to upgrade to.",
"warn")
yield commands.CloseConnection(self.context.client)
self.client_state = self.server_state = self.state_errored
return
if self.debug:
yield commands.Log(
f"{self.debug}[http] upgrading to {self.child_layer}",
"debug")
yield from self.child_layer.handle_event(events.Start())
self._handle_event = self.passthrough
return
def test_receive_data_before_server_connected(tctx):
"""
assert that data received before a server connection is established
will still be forwarded.
"""
assert (Playbook(tcp.TCPLayer(tctx), hooks=False) << OpenConnection(
tctx.server) >> DataReceived(tctx.client, b"hello!") >> reply(
None, to=-2) << SendData(tctx.server, b"hello!"))
def no_flow_hooks():
assert (
Playbook(tcp.TCPLayer(tctx, ignore=ignore), hooks=True)
<< OpenConnection(tctx.server)
>> reply(None)
>> DataReceived(tctx.client, b"hello!")
<< SendData(tctx.server, b"hello!")
)
def test_receive_data_after_half_close(tctx):
"""
data received after the other connection has been half-closed should still be forwarded.
"""
assert (Playbook(tcp.TCPLayer(tctx), hooks=False) <<
OpenConnection(tctx.server) >> reply(None) >> DataReceived(
tctx.client, b"eof-delimited-request") << SendData(
tctx.server, b"eof-delimited-request") >> ConnectionClosed(
tctx.client) <<
CloseConnection(tctx.server, half_close=True) >> DataReceived(
tctx.server, b"i'm late") << SendData(tctx.client, b"i'm late")
>> ConnectionClosed(tctx.server) << CloseConnection(tctx.client))
def test_open_connection_err(tctx):
f = Placeholder(TCPFlow)
assert (
Playbook(tcp.TCPLayer(tctx))
<< tcp.TcpStartHook(f)
>> reply()
<< OpenConnection(tctx.server)
>> reply("Connect call failed")
<< tcp.TcpErrorHook(f)
>> reply()
<< CloseConnection(tctx.client)
)
def receive_handshake_data(
self,
data: bytes) -> layer.CommandGenerator[Tuple[bool, Optional[str]]]:
if self.client_hello_parsed:
return (yield from super().receive_handshake_data(data))
self.recv_buffer.extend(data)
try:
client_hello = parse_client_hello(self.recv_buffer)
except ValueError:
return False, f"Cannot parse ClientHello: {self.recv_buffer.hex()}"
if client_hello:
self.client_hello_parsed = True
else:
return False, None
self.conn.sni = client_hello.sni
self.conn.alpn_offers = client_hello.alpn_protocols
tls_clienthello = ClientHelloData(self.context, client_hello)
yield TlsClienthelloHook(tls_clienthello)
if tls_clienthello.ignore_connection:
# we've figured out that we don't want to intercept this connection, so we assign fake connection objects
# to all TLS layers. This makes the real connection contents just go through.
self.conn = self.tunnel_connection = connection.Client(
("ignore-conn", 0), ("ignore-conn", 0), time.time())
parent_layer = self.context.layers[self.context.layers.index(self)
- 1]
if isinstance(parent_layer, ServerTLSLayer):
parent_layer.conn = parent_layer.tunnel_connection = connection.Server(
None)
self.child_layer = tcp.TCPLayer(self.context, ignore=True)
yield from self.event_to_child(
events.DataReceived(self.context.client,
bytes(self.recv_buffer)))
self.recv_buffer.clear()
return True, None
if tls_clienthello.establish_server_tls_first and not self.context.server.tls_established:
err = yield from self.start_server_tls()
if err:
yield commands.Log(
f"Unable to establish TLS connection with server ({err}). "
f"Trying to establish TLS with client anyway.")
yield from self.start_tls()
if not self.conn.connected:
return False, "connection closed early"
ret = yield from super().receive_handshake_data(bytes(
self.recv_buffer))
self.recv_buffer.clear()
return ret
def test_inject(tctx):
"""inject data into an open connection."""
f = Placeholder(TCPFlow)
assert (
Playbook(tcp.TCPLayer(tctx)) << tcp.TcpStartHook(f) >>
TcpMessageInjected(f, TCPMessage(True, b"hello!")) >> reply(to=-2) <<
OpenConnection(tctx.server) >> reply(None) << tcp.TcpMessageHook(f) >>
reply() << SendData(tctx.server, b"hello!")
# and the other way...
>> TcpMessageInjected(
f, TCPMessage(False, b"I have already done the greeting for you."))
<< tcp.TcpMessageHook(f) >> reply() << SendData(
tctx.client, b"I have already done the greeting for you.") << None)
assert len(f().messages) == 2
