[sans-io] add mode tests, minor fixes

This commit is contained in:
Maximilian Hils 2020-12-12 21:22:48 +01:00
parent 4c75765387
commit 641b1c425b
9 changed files with 201 additions and 37 deletions

View File

@ -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())

View File

@ -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})"

View File

@ -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

View File

@ -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):

View File

@ -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()

View File

@ -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()

View File

@ -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

View File

@ -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

View File

@ -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)
)