def test_tunnel_handshake_start(tctx: Context, success):
server = Server(("proxy", 1234))
server.state = ConnectionState.OPEN
tl = TTunnelLayer(tctx, server, tctx.server)
tl.child_layer = TChildLayer(tctx)
assert repr(tl)
playbook = Playbook(tl, logs=True)
(playbook << SendData(server, b"handshake-hello") >> DataReceived(
tctx.client, b"client-hello") >> DataReceived(
server, b"handshake-" + success.encode()) << SendData(
server, b"handshake-" + success.encode()))
if success == "success":
playbook << Log("Got start. Server state: OPEN")
else:
playbook << CloseConnection(server)
playbook << Log("Got start. Server state: CLOSED")
playbook << SendData(tctx.client, b"client-hello-reply")
if success == "success":
playbook >> DataReceived(server, b"tunneled-server-hello")
playbook << SendData(server, b"tunneled-server-hello-reply")
assert playbook
def test_tunnel_handshake_command(tctx: Context, success):
server = Server(("proxy", 1234))
tl = TTunnelLayer(tctx, server, tctx.server)
tl.child_layer = TChildLayer(tctx)
playbook = Playbook(tl, logs=True)
(playbook << Log("Got start. Server state: CLOSED") >> DataReceived(
tctx.client, b"client-hello") << SendData(
tctx.client, b"client-hello-reply") >> DataReceived(
tctx.client, b"open") << OpenConnection(server) >> reply(None)
<< SendData(server, b"handshake-hello") >> DataReceived(
server, b"handshake-" + success.encode()) << SendData(
server, b"handshake-" + success.encode()))
if success == "success":
assert (playbook << Log(f"Opened: err=None. Server state: OPEN") >>
DataReceived(server, b"tunneled-server-hello") << SendData(
server, b"tunneled-server-hello-reply") >>
ConnectionClosed(tctx.client) << Log("Got client close.") <<
CloseConnection(tctx.client))
assert tl.tunnel_state is tunnel.TunnelState.OPEN
assert (playbook >> ConnectionClosed(server) <<
Log("Got server close.") << CloseConnection(server))
assert tl.tunnel_state is tunnel.TunnelState.CLOSED
else:
assert (playbook << CloseConnection(server) <<
Log("Opened: err='handshake error'. Server state: CLOSED"))
assert tl.tunnel_state is tunnel.TunnelState.CLOSED
def test_disconnect_during_handshake_command(tctx: Context, disconnect):
server = Server(("proxy", 1234))
tl = TTunnelLayer(tctx, server, tctx.server)
tl.child_layer = TChildLayer(tctx)
playbook = Playbook(tl, logs=True)
assert (
playbook
<< Log("Got start. Server state: CLOSED")
>> DataReceived(tctx.client, b"client-hello")
<< SendData(tctx.client, b"client-hello-reply")
>> DataReceived(tctx.client, b"open")
<< OpenConnection(server)
>> reply(None)
<< SendData(server, b"handshake-hello")
)
if disconnect == "client":
assert (
playbook
>> ConnectionClosed(tctx.client)
>> ConnectionClosed(server) # proxyserver will cancel all other connections as well.
<< CloseConnection(server)
<< Log("Opened: err='connection closed without notice'. Server state: CLOSED")
<< Log("Got client close.")
<< CloseConnection(tctx.client)
)
else:
assert (
playbook
>> ConnectionClosed(server)
<< CloseConnection(server)
<< Log("Opened: err='connection closed without notice'. Server state: CLOSED")
)
def test_disconnect_during_handshake_start(tctx: Context, disconnect):
server = Server(("proxy", 1234))
server.state = ConnectionState.OPEN
tl = TTunnelLayer(tctx, server, tctx.server)
tl.child_layer = TChildLayer(tctx)
playbook = Playbook(tl, logs=True)
assert (
playbook
<< SendData(server, b"handshake-hello")
)
if disconnect == "client":
assert (
playbook
>> ConnectionClosed(tctx.client)
>> ConnectionClosed(server) # proxyserver will cancel all other connections as well.
<< CloseConnection(server)
<< Log("Got start. Server state: CLOSED")
<< Log("Got client close.")
<< CloseConnection(tctx.client)
)
else:
assert (
playbook
>> ConnectionClosed(server)
<< CloseConnection(server)
<< Log("Got start. Server state: CLOSED")
)
def test_tunnel_default_impls(tctx: Context):
"""
Some tunnels don't need certain features, so the default behaviour
should be to be transparent.
"""
server = Server(None)
server.state = ConnectionState.OPEN
tl = tunnel.TunnelLayer(tctx, server, tctx.server)
tl.child_layer = TChildLayer(tctx)
playbook = Playbook(tl, logs=True)
assert (
playbook
<< Log("Got start. Server state: OPEN")
>> DataReceived(server, b"server-hello")
<< SendData(server, b"server-hello-reply")
)
assert tl.tunnel_state is tunnel.TunnelState.OPEN
assert (
playbook
>> ConnectionClosed(server)
<< Log("Got server close.")
<< CloseConnection(server)
)
assert tl.tunnel_state is tunnel.TunnelState.CLOSED
assert (
playbook
>> DataReceived(tctx.client, b"open")
<< OpenConnection(server)
>> reply(None)
<< Log("Opened: err=None. Server state: OPEN")
>> DataReceived(server, b"half-close")
<< CloseConnection(server, half_close=True)
)
def test_tunnel_openconnection_error(tctx: Context):
server = Server(("proxy", 1234))
tl = TTunnelLayer(tctx, server, tctx.server)
tl.child_layer = TChildLayer(tctx)
playbook = Playbook(tl, logs=True)
assert (playbook << Log("Got start. Server state: CLOSED") >> DataReceived(
tctx.client, b"open") << OpenConnection(server))
assert tl.tunnel_state is tunnel.TunnelState.ESTABLISHING
assert (playbook >> reply("IPoAC packet dropped.") <<
Log("Opened: err='IPoAC packet dropped.'. Server state: CLOSED"))
assert tl.tunnel_state is tunnel.TunnelState.CLOSED
def test_ping(ws_testdata):
tctx, playbook, flow = ws_testdata
assert (
playbook
<< websocket.WebsocketStartHook(flow)
>> reply()
>> DataReceived(tctx.client, masked_bytes(b"\x89\x11ping-with-payload"))
<< Log("Received WebSocket ping from client (payload: b'ping-with-payload')")
<< SendData(tctx.server, masked(b"\x89\x11ping-with-payload"))
>> DataReceived(tctx.server, b"\x8a\x11pong-with-payload")
<< Log("Received WebSocket pong from server (payload: b'pong-with-payload')")
<< SendData(tctx.client, b"\x8a\x11pong-with-payload")
)
assert not flow.websocket.messages
def test_socks5_err(data: bytes, err: bytes, msg: str, tctx: Context):
playbook = (Playbook(modes.Socks5Proxy(tctx), logs=True) >> DataReceived(
tctx.client, data))
if err:
playbook << SendData(tctx.client, err)
playbook << CloseConnection(tctx.client)
playbook << Log(msg)
assert playbook
def _handle_event(self, event: Event) -> layer.CommandGenerator[None]:
if isinstance(event, Start):
yield Log(f"Got start. Server state: {self.context.server.state.name}")
elif isinstance(event, DataReceived) and event.data == b"client-hello":
yield SendData(self.context.client, b"client-hello-reply")
elif isinstance(event, DataReceived) and event.data == b"server-hello":
yield SendData(self.context.server, b"server-hello-reply")
elif isinstance(event, DataReceived) and event.data == b"open":
err = yield OpenConnection(self.context.server)
yield Log(f"Opened: {err=}. Server state: {self.context.server.state.name}")
elif isinstance(event, DataReceived) and event.data == b"half-close":
err = yield CloseConnection(event.connection, half_close=True)
elif isinstance(event, ConnectionClosed):
yield Log(f"Got {event.connection.__class__.__name__.lower()} close.")
yield CloseConnection(event.connection)
else:
raise AssertionError
def test_socks5_premature_close(tctx: Context):
assert (
Playbook(modes.Socks5Proxy(tctx), logs=True)
>> DataReceived(tctx.client, b"\x05")
>> ConnectionClosed(tctx.client)
<< Log(r"Client closed connection before completing SOCKS5 handshake: b'\x05'")
<< CloseConnection(tctx.client)
)
def test_unknown_ext(ws_testdata):
tctx, playbook, flow = ws_testdata
flow.response.headers["Sec-WebSocket-Extensions"] = "funky-bits; param=42"
assert (
playbook
<< Log("Ignoring unknown WebSocket extension 'funky-bits'.")
<< websocket.WebsocketStartHook(flow)
>> reply()
)
def test_unknown_ext(ws_testdata):
tctx, playbook = ws_testdata
flow = Placeholder(WebSocketFlow)
# noinspection PyUnresolvedReferences
http_flow: HTTPFlow = playbook.layer.flow.handshake_flow
http_flow.response.headers[
"Sec-WebSocket-Extensions"] = "funky-bits; param=42"
assert (
playbook << Log("Ignoring unknown WebSocket extension 'funky-bits'.")
<< websocket.WebsocketStartHook(flow) >> reply())
def test_transparent_failure(tctx: Context, monkeypatch):
"""Test that we recover from a transparent mode resolve error."""
def raise_err(sock):
raise RuntimeError("platform-specific error")
monkeypatch.setattr(platform, "original_addr", raise_err)
assert (Playbook(modes.TransparentProxy(tctx), logs=True) << GetSocket(
tctx.client
) >> reply(object()) << Log(
"Transparent mode failure: RuntimeError('platform-specific error')",
"info"))
def test_upgrade(tctx, proto):
"""Test a HTTP -> WebSocket upgrade with different protocols enabled"""
if proto != "websocket":
tctx.options.websocket = False
if proto != "tcp":
tctx.options.rawtcp = False
tctx.server.address = ("example.com", 80)
tctx.server.state = ConnectionState.OPEN
http_flow = Placeholder(HTTPFlow)
playbook = Playbook(http.HttpLayer(tctx, HTTPMode.transparent))
(
playbook
>> DataReceived(tctx.client,
b"GET / HTTP/1.1\r\n"
b"Connection: upgrade\r\n"
b"Upgrade: websocket\r\n"
b"Sec-WebSocket-Version: 13\r\n"
b"\r\n")
<< http.HttpRequestHeadersHook(http_flow)
>> reply()
<< http.HttpRequestHook(http_flow)
>> reply()
<< SendData(tctx.server, b"GET / HTTP/1.1\r\n"
b"Connection: upgrade\r\n"
b"Upgrade: websocket\r\n"
b"Sec-WebSocket-Version: 13\r\n"
b"\r\n")
>> DataReceived(tctx.server, b"HTTP/1.1 101 Switching Protocols\r\n"
b"Upgrade: websocket\r\n"
b"Connection: Upgrade\r\n"
b"\r\n")
<< http.HttpResponseHeadersHook(http_flow)
>> reply()
<< http.HttpResponseHook(http_flow)
>> reply()
<< SendData(tctx.client, b"HTTP/1.1 101 Switching Protocols\r\n"
b"Upgrade: websocket\r\n"
b"Connection: Upgrade\r\n"
b"\r\n")
)
if proto == "websocket":
assert playbook << WebsocketStartHook(Placeholder(WebSocketFlow))
elif proto == "tcp":
assert playbook << TcpStartHook(Placeholder(TCPFlow))
else:
assert (
playbook
<< Log("Sent HTTP 101 response, but no protocol is enabled to upgrade to.", "warn")
<< CloseConnection(tctx.client)
)
def test_no_normalization(tctx, normalize):
"""Test that we don't normalize headers when we just pass them through."""
tctx.options.normalize_outbound_headers = normalize
tctx.options.validate_inbound_headers = False
server = Placeholder(Server)
flow = Placeholder(HTTPFlow)
playbook, cff = start_h2_client(tctx)
request_headers = list(example_request_headers) + [(b"Should-Not-Be-Capitalized! ", b" :) ")]
request_headers_lower = [(k.lower(), v) for (k, v) in request_headers]
response_headers = list(example_response_headers) + [(b"Same", b"Here")]
response_headers_lower = [(k.lower(), v) for (k, v) in response_headers]
initial = Placeholder(bytes)
assert (
playbook
>> DataReceived(tctx.client,
cff.build_headers_frame(request_headers, flags=["END_STREAM"]).serialize())
<< http.HttpRequestHeadersHook(flow)
>> reply()
<< http.HttpRequestHook(flow)
>> reply()
<< OpenConnection(server)
>> reply(None, side_effect=make_h2)
<< SendData(server, initial)
)
frames = decode_frames(initial())
assert [type(x) for x in frames] == [
hyperframe.frame.SettingsFrame,
hyperframe.frame.HeadersFrame,
]
assert hpack.hpack.Decoder().decode(frames[1].data, True) == request_headers_lower if normalize else request_headers
sff = FrameFactory()
(
playbook
>> DataReceived(server, sff.build_headers_frame(response_headers, flags=["END_STREAM"]).serialize())
<< http.HttpResponseHeadersHook(flow)
>> reply()
<< http.HttpResponseHook(flow)
>> reply()
)
if normalize:
playbook << Log("Lowercased 'Same' header as uppercase is not allowed with HTTP/2.")
hdrs = response_headers_lower if normalize else response_headers
assert playbook << SendData(tctx.client, cff.build_headers_frame(hdrs, flags=["END_STREAM"]).serialize())
assert flow().request.headers.fields == ((b"Should-Not-Be-Capitalized! ", b" :) "),)
assert flow().response.headers.fields == ((b"Same", b"Here"),)
def test_socks5_auth_fail(client_greeting: bytes, server_choice: bytes,
client_auth: bytes, err: bytes, msg: str,
tctx: Context):
ProxyAuth().load(tctx.options)
tctx.options.proxyauth = "user:password"
playbook = (Playbook(modes.Socks5Proxy(tctx), logs=True) >> DataReceived(
tctx.client, client_greeting))
if server_choice is None:
playbook << SendData(tctx.client, err)
else:
playbook << SendData(tctx.client, server_choice)
playbook >> DataReceived(tctx.client, client_auth)
playbook << modes.Socks5AuthHook(Placeholder(modes.Socks5AuthData))
playbook >> reply()
playbook << SendData(tctx.client, err)
playbook << CloseConnection(tctx.client)
playbook << Log(msg)
assert playbook
