From 641b1c425bc742cd5e07f93da536ce12ca8972ca Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Sat, 12 Dec 2020 21:22:48 +0100 Subject: [PATCH] [sans-io] add mode tests, minor fixes --- mitmproxy/addons/proxyserver.py | 4 +- mitmproxy/proxy2/commands.py | 2 +- mitmproxy/proxy2/context.py | 8 +- mitmproxy/proxy2/events.py | 3 - mitmproxy/proxy2/layers/modes.py | 62 ++++++--- mitmproxy/proxy2/server.py | 4 +- mitmproxy/proxy2/utils.py | 4 +- setup.cfg | 2 + test/mitmproxy/proxy2/layers/test_modes.py | 149 ++++++++++++++++++++- 9 files changed, 201 insertions(+), 37 deletions(-) diff --git a/mitmproxy/addons/proxyserver.py b/mitmproxy/addons/proxyserver.py index ab8e23aaa..e8f87166a 100644 --- a/mitmproxy/addons/proxyserver.py +++ b/mitmproxy/addons/proxyserver.py @@ -2,7 +2,7 @@ import asyncio import warnings from typing import Optional -from mitmproxy import controller, ctx, eventsequence, flow, log, master, options +from mitmproxy import controller, ctx, eventsequence, flow, log, master, options, platform from mitmproxy.flow import Error from mitmproxy.proxy2 import commands from mitmproxy.proxy2 import server @@ -93,6 +93,8 @@ class Proxyserver: def configure(self, updated): if not self.is_running: return + if "mode" in updated and ctx.options.mode == "transparent": # pragma: no cover + platform.init_transparent_mode() if any(x in updated for x in ["server", "listen_host", "listen_port"]): asyncio.create_task(self.refresh_server()) diff --git a/mitmproxy/proxy2/commands.py b/mitmproxy/proxy2/commands.py index 3f9af7e3c..25deca8be 100644 --- a/mitmproxy/proxy2/commands.py +++ b/mitmproxy/proxy2/commands.py @@ -145,4 +145,4 @@ class Log(Command): self.level = level def __repr__(self): - return f"Log({self.message}, {self.level})" + return f"Log({self.message!r}, {self.level!r})" diff --git a/mitmproxy/proxy2/context.py b/mitmproxy/proxy2/context.py index f6f0d60ac..60edff1c6 100644 --- a/mitmproxy/proxy2/context.py +++ b/mitmproxy/proxy2/context.py @@ -102,7 +102,7 @@ class Connection(serializable.Serializable, metaclass=ABCMeta): class Client(Connection): - state = ConnectionState.OPEN + state: ConnectionState = ConnectionState.OPEN peername: Address sockname: Address @@ -205,9 +205,9 @@ class Client(Connection): class Server(Connection): - state = ConnectionState.CLOSED - peername = None - sockname = None + state: ConnectionState = ConnectionState.CLOSED + peername: Optional[Address] = None + sockname: Optional[Address] = None address: Optional[Address] timestamp_start: Optional[float] = None diff --git a/mitmproxy/proxy2/events.py b/mitmproxy/proxy2/events.py index c1875d346..fe3077787 100644 --- a/mitmproxy/proxy2/events.py +++ b/mitmproxy/proxy2/events.py @@ -94,9 +94,6 @@ class OpenConnectionReply(CommandReply): reply: typing.Optional[str] """error message""" - def __repr__(self): - return super().__repr__() - @dataclass(repr=False) class HookReply(CommandReply): diff --git a/mitmproxy/proxy2/layers/modes.py b/mitmproxy/proxy2/layers/modes.py index 839eddac7..3f8c64f89 100644 --- a/mitmproxy/proxy2/layers/modes.py +++ b/mitmproxy/proxy2/layers/modes.py @@ -1,27 +1,12 @@ +from abc import ABCMeta + from mitmproxy import platform from mitmproxy.net import server_spec from mitmproxy.proxy2 import commands, events, layer -from mitmproxy.proxy2.context import Server from mitmproxy.proxy2.layers import tls from mitmproxy.proxy2.utils import expect -class ReverseProxy(layer.Layer): - @expect(events.Start) - def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]: - spec = server_spec.parse_with_mode(self.context.options.mode)[1] - self.context.server = Server(spec.address) - child_layer: layer.Layer - if spec.scheme not in ("http", "tcp"): - if not self.context.options.keep_host_header: - self.context.server.sni = spec.address[0].encode() - child_layer = tls.ServerTLSLayer(self.context) - else: - child_layer = layer.NextLayer(self.context) - self._handle_event = child_layer.handle_event - yield from child_layer.handle_event(event) - - class HttpProxy(layer.Layer): @expect(events.Start) def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]: @@ -30,7 +15,42 @@ class HttpProxy(layer.Layer): yield from child_layer.handle_event(event) -class TransparentProxy(layer.Layer): +class DestinationKnown(layer.Layer, metaclass=ABCMeta): + child_layer: layer.Layer + + def finish_start(self): + if self.context.options.connection_strategy == "eager": + err = yield commands.OpenConnection(self.context.server) + if err: + yield commands.CloseConnection(self.context.client) + self._handle_event = self.done + return + + self._handle_event = self.child_layer.handle_event + yield from self.child_layer.handle_event(events.Start()) + + @expect(events.DataReceived, events.ConnectionClosed) + def done(self, _) -> layer.CommandGenerator[None]: + yield from () + + +class ReverseProxy(DestinationKnown): + @expect(events.Start) + def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]: + spec = server_spec.parse_with_mode(self.context.options.mode)[1] + self.context.server.address = spec.address + + if spec.scheme not in ("http", "tcp"): + if not self.context.options.keep_host_header: + self.context.server.sni = spec.address[0].encode() + self.child_layer = tls.ServerTLSLayer(self.context) + else: + self.child_layer = layer.NextLayer(self.context) + + yield from self.finish_start() + + +class TransparentProxy(DestinationKnown): @expect(events.Start) def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]: assert platform.original_addr is not None @@ -40,6 +60,6 @@ class TransparentProxy(layer.Layer): except Exception as e: yield commands.Log(f"Transparent mode failure: {e!r}") - child_layer = layer.NextLayer(self.context) - self._handle_event = child_layer.handle_event - yield from child_layer.handle_event(event) + self.child_layer = layer.NextLayer(self.context) + + yield from self.finish_start() diff --git a/mitmproxy/proxy2/server.py b/mitmproxy/proxy2/server.py index e20a62fad..a65ab5402 100644 --- a/mitmproxy/proxy2/server.py +++ b/mitmproxy/proxy2/server.py @@ -343,7 +343,7 @@ class StreamConnectionHandler(ConnectionHandler, metaclass=abc.ABCMeta): self.transports[client] = ConnectionIO(handler=None, reader=reader, writer=writer) -class SimpleConnectionHandler(StreamConnectionHandler): +class SimpleConnectionHandler(StreamConnectionHandler): # pragma: no cover """Simple handler that does not really process any hooks.""" hook_handlers: typing.Dict[str, typing.Callable] @@ -364,7 +364,7 @@ class SimpleConnectionHandler(StreamConnectionHandler): pass # print(message, file=sys.stderr if level in ("error", "warn") else sys.stdout) -if __name__ == "__main__": +if __name__ == "__main__": # pragma: no cover # simple standalone implementation for testing. loop = asyncio.get_event_loop() diff --git a/mitmproxy/proxy2/utils.py b/mitmproxy/proxy2/utils.py index 6426d62e4..9ed9d0fba 100644 --- a/mitmproxy/proxy2/utils.py +++ b/mitmproxy/proxy2/utils.py @@ -17,7 +17,7 @@ def expect(*event_types): @functools.wraps(f) def _check_event_type(self, event: events.Event): if isinstance(event, event_types): - yield from f(self, event) + return f(self, event) else: event_types_str = '|'.join(e.__name__ for e in event_types) or "no events" raise AssertionError( @@ -26,7 +26,7 @@ def expect(*event_types): ) return _check_event_type - else: + else: # pragma: no cover return f return decorator diff --git a/setup.cfg b/setup.cfg index 554d29b19..9e00816bb 100644 --- a/setup.cfg +++ b/setup.cfg @@ -94,6 +94,8 @@ exclude = mitmproxy/proxy/protocol/tls.py mitmproxy/proxy/root_context.py mitmproxy/proxy/server.py + mitmproxy/proxy2/server.py + mitmproxy/proxy2/layers/tls.py mitmproxy/utils/bits.py pathod/language/actions.py pathod/language/base.py diff --git a/test/mitmproxy/proxy2/layers/test_modes.py b/test/mitmproxy/proxy2/layers/test_modes.py index 3dc7accf5..500082172 100644 --- a/test/mitmproxy/proxy2/layers/test_modes.py +++ b/test/mitmproxy/proxy2/layers/test_modes.py @@ -2,13 +2,16 @@ import copy import pytest +from mitmproxy import platform from mitmproxy.proxy.protocol.http import HTTPMode -from mitmproxy.proxy2.commands import CloseConnection, OpenConnection, SendData +from mitmproxy.proxy2.commands import CloseConnection, OpenConnection, SendData, GetSocket, Log from mitmproxy.proxy2.context import Client, Context, Server -from mitmproxy.proxy2.events import DataReceived +from mitmproxy.proxy2.events import DataReceived, ConnectionClosed from mitmproxy.proxy2.layer import NextLayer, NextLayerHook -from mitmproxy.proxy2.layers import http, modes +from mitmproxy.proxy2.layers import http, modes, tcp, tls +from mitmproxy.proxy2.layers.tcp import TcpStartHook, TcpMessageHook from mitmproxy.proxy2.layers.tls import ClientTLSLayer, TlsStartHook +from mitmproxy.tcp import TCPFlow from test.mitmproxy.proxy2.layers.test_tls import reply_tls_start from test.mitmproxy.proxy2.tutils import Placeholder, Playbook, reply, reply_next_layer @@ -130,3 +133,143 @@ def test_reverse_proxy(tctx, keep_host_header): << SendData(tctx.client, b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") ) assert server().address == ("localhost", 8000) + + +@pytest.mark.parametrize("patch", [True, False]) +@pytest.mark.parametrize("connection_strategy", ["eager", "lazy"]) +def test_reverse_proxy_tcp_over_tls(tctx: Context, monkeypatch, patch, connection_strategy): + """ + Test + client --TCP-- mitmproxy --TCP over TLS-- server + reverse proxying. + """ + + if patch: + monkeypatch.setattr(tls, "ServerTLSLayer", tls.MockTLSLayer) + + flow = Placeholder(TCPFlow) + data = Placeholder(bytes) + tctx.options.mode = "reverse:https://localhost:8000" + tctx.options.connection_strategy = connection_strategy + playbook = Playbook(modes.ReverseProxy(tctx)) + if connection_strategy == "eager": + ( + playbook + << OpenConnection(tctx.server) + >> DataReceived(tctx.client, b"\x01\x02\x03") + >> reply(None, to=OpenConnection(tctx.server)) + ) + else: + ( + playbook + >> DataReceived(tctx.client, b"\x01\x02\x03") + ) + if patch: + ( + playbook + << NextLayerHook(Placeholder(NextLayer)) + >> reply_next_layer(tcp.TCPLayer) + << TcpStartHook(flow) + >> reply() + ) + if connection_strategy == "lazy": + ( + playbook + << OpenConnection(tctx.server) + >> reply(None) + ) + assert ( + playbook + << TcpMessageHook(flow) + >> reply() + << SendData(tctx.server, data) + ) + assert data() == b"\x01\x02\x03" + else: + if connection_strategy == "lazy": + ( + playbook + << NextLayerHook(Placeholder(NextLayer)) + >> reply_next_layer(tcp.TCPLayer) + << TcpStartHook(flow) + >> reply() + << OpenConnection(tctx.server) + >> reply(None) + ) + assert ( + playbook + << TlsStartHook(Placeholder()) + >> reply_tls_start() + << SendData(tctx.server, data) + ) + assert tls.parse_client_hello(data()).sni == b"localhost" + + +@pytest.mark.parametrize("connection_strategy", ["eager", "lazy"]) +def test_transparent_tcp(tctx: Context, monkeypatch, connection_strategy): + monkeypatch.setattr(platform, "original_addr", lambda sock: ("address", 22)) + + flow = Placeholder(TCPFlow) + tctx.options.connection_strategy = connection_strategy + + sock = object() + playbook = Playbook(modes.TransparentProxy(tctx)) + ( + playbook + << GetSocket(tctx.client) + >> reply(sock) + ) + if connection_strategy == "lazy": + assert playbook + else: + assert ( + playbook + << OpenConnection(tctx.server) + >> reply(None) + >> DataReceived(tctx.server, b"hello") + << NextLayerHook(Placeholder(NextLayer)) + >> reply_next_layer(tcp.TCPLayer) + << TcpStartHook(flow) + >> reply() + << TcpMessageHook(flow) + >> reply() + << SendData(tctx.client, b"hello") + ) + assert flow().messages[0].content == b"hello" + assert not flow().messages[0].from_client + + assert tctx.server.address == ("address", 22) + + +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_reverse_eager_connect_failure(tctx: Context): + """ + Test + client --TCP-- mitmproxy --TCP over TLS-- server + reverse proxying. + """ + + tctx.options.mode = "reverse:https://localhost:8000" + tctx.options.connection_strategy = "eager" + playbook = Playbook(modes.ReverseProxy(tctx)) + assert ( + playbook + << OpenConnection(tctx.server) + >> reply("IPoAC unstable") + << CloseConnection(tctx.client) + >> ConnectionClosed(tctx.client) + )