mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-25 09:37:37 +00:00
[sans-io] add mode tests, minor fixes
This commit is contained in:
parent
4c75765387
commit
641b1c425b
@ -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())
|
||||
|
||||
|
@ -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})"
|
||||
|
@ -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
|
||||
|
@ -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):
|
||||
|
@ -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()
|
||||
|
@ -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()
|
||||
|
||||
|
@ -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
|
||||
|
@ -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
|
||||
|
@ -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)
|
||||
)
|
||||
|
Loading…
Reference in New Issue
Block a user