From edaf3219fc3d78f0a7b2b72a1102e8adaecaf6b6 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Thu, 26 Oct 2017 18:32:03 +0200 Subject: [PATCH] [sans-io] improve logging --- mitmproxy/addons/proxyserver.py | 15 +++++++-------- mitmproxy/proxy2/commands.py | 19 +++++++++++-------- mitmproxy/proxy2/context.py | 6 ------ mitmproxy/proxy2/layer.py | 3 ++- mitmproxy/proxy2/layers/http.py | 10 +++++----- mitmproxy/proxy2/layers/tls.py | 4 ++-- mitmproxy/proxy2/server.py | 33 ++++++++++++++++++++------------- 7 files changed, 47 insertions(+), 43 deletions(-) diff --git a/mitmproxy/addons/proxyserver.py b/mitmproxy/addons/proxyserver.py index 6770282d2..d2008768f 100644 --- a/mitmproxy/addons/proxyserver.py +++ b/mitmproxy/addons/proxyserver.py @@ -42,8 +42,8 @@ class ProxyConnectionHandler(server.ConnectionHandler): if hook.blocking: self.server_event(events.HookReply(hook)) - def _debug(self, *args): - x = log.LogEntry(" ".join(str(x) for x in args), "warn") + def log(self, message: str, level: str = "info") -> None: + x = log.LogEntry(message, level) x.reply = controller.DummyReply() self.event_queue.put(("log", x)) @@ -60,6 +60,7 @@ class Proxyserver: self.event_queue = None self.options = None self._lock = asyncio.Lock() + self.is_running = False def load(self, loader): loader.add_option( @@ -72,6 +73,8 @@ class Proxyserver: self.options = ctx.options self.event_queue = ctx.master.event_queue threading.Thread(target=self.loop.run_forever, daemon=True).start() + self.is_running = True + self.configure(["listen_port"]) async def start(self): async with self._lock: @@ -98,14 +101,10 @@ class Proxyserver: ).handle_client() def configure(self, updated): + if not self.is_running: + return if "listen_port" in updated: self.listen_port = ctx.options.listen_port + 1 # not sure if this actually required... self.loop.call_soon_threadsafe(lambda: asyncio.ensure_future(self.start())) - - def request(self, flow): - pass - # test live options changes. - # print("Changing port...") - # ctx.options.listen_port += 1 diff --git a/mitmproxy/proxy2/commands.py b/mitmproxy/proxy2/commands.py index d9577cec3..9434c423d 100644 --- a/mitmproxy/proxy2/commands.py +++ b/mitmproxy/proxy2/commands.py @@ -23,7 +23,8 @@ class Command: Example: - reply = yield Hook("requestheaders", flow) + reply = yield Hook("requestheaders", flow) # blocking command + yield Log("hello world", "info") # non-blocking """ def __repr__(self): @@ -84,17 +85,19 @@ class Hook(Command): self.data = data -class Log(Hook): +class Log(Command): blocking = False + message: str + level: str - # this is more of a hack at the moment. - def __init__(self, *args, level="info"): - if len(args) == 1: - args = args[0] - super().__init__("log", log.LogEntry(str(args), level)) + def __init__(self, message, level="info"): + assert isinstance(message, str) + assert isinstance(level, str) + self.message = message + self.level = level def __repr__(self): - return f"Log: {self.data.msg}" + return f"Log({self.message}, {self.level})" TCommandGenerator = typing.Generator[Command, typing.Any, None] diff --git a/mitmproxy/proxy2/context.py b/mitmproxy/proxy2/context.py index 838ce8f54..86bf907d1 100644 --- a/mitmproxy/proxy2/context.py +++ b/mitmproxy/proxy2/context.py @@ -35,12 +35,6 @@ class Server(Connection): class Context: """ Layers get a context object that has all contextual information they require. - For now, the only required property is the client connection, with ClientServerContext - adding the server connection. - - TODO: We may just use a simple context class that has _all_ attributes we ever think of? - Alternatively, we could have a `.get(attr)`, that mimicks what the transparent attribute - lookup did in the previous implementation. """ client: Client diff --git a/mitmproxy/proxy2/layer.py b/mitmproxy/proxy2/layer.py index d4c8a3e4b..142cb77f1 100644 --- a/mitmproxy/proxy2/layer.py +++ b/mitmproxy/proxy2/layer.py @@ -20,12 +20,13 @@ class Paused(typing.NamedTuple): class Layer: context: Context _paused: typing.Optional[Paused] + _paused_event_queue: typing.Deque[events.Event] def __init__(self, context: Context) -> None: self.context = context self.context.layers.append(self) self._paused = None - self._paused_event_queue: typing.Deque[events.Event] = collections.deque() + self._paused_event_queue = collections.deque() def _debug(self, *args): pass # print(*args) diff --git a/mitmproxy/proxy2/layers/http.py b/mitmproxy/proxy2/layers/http.py index 55145a637..429d122ea 100644 --- a/mitmproxy/proxy2/layers/http.py +++ b/mitmproxy/proxy2/layers/http.py @@ -62,7 +62,7 @@ class HTTPLayer(Layer): if event is h11.NEED_DATA: return elif isinstance(event, h11.Request): - yield commands.Log("requestheaders", event) + yield commands.Log(f"requestheaders: {event}") if self.client_conn.client_is_waiting_for_100_continue: raise NotImplementedError() @@ -83,7 +83,7 @@ class HTTPLayer(Layer): self.flow_events[0].append(event) elif isinstance(event, h11.EndOfMessage): self.flow_events[0].append(event) - yield commands.Log("request", self.flow_events) + yield commands.Log(f"request {self.flow_events}") yield from self._send_request() return else: @@ -93,7 +93,7 @@ class HTTPLayer(Layer): if not self.context.server.connected: err = yield commands.OpenConnection(self.context.server) if err: - yield commands.Log("error", err) + yield commands.Log(f"error {err}") yield commands.CloseConnection(self.context.client) self._handle_event = self.done return @@ -107,7 +107,7 @@ class HTTPLayer(Layer): if event is h11.NEED_DATA: return elif isinstance(event, h11.Response): - yield commands.Log("responseheaders", event) + yield commands.Log(f"responseheaders {event}") self.flow_events[1].append(event) self.state = self.read_response_body @@ -131,7 +131,7 @@ class HTTPLayer(Layer): self.flow_events[1].append(event) elif isinstance(event, h11.EndOfMessage): self.flow_events[1].append(event) - yield commands.Log("response", self.flow_events) + yield commands.Log(f"response {self.flow_events}") yield from self._send_response() return else: diff --git a/mitmproxy/proxy2/layers/tls.py b/mitmproxy/proxy2/layers/tls.py index 1f44cf666..925a864ea 100644 --- a/mitmproxy/proxy2/layers/tls.py +++ b/mitmproxy/proxy2/layers/tls.py @@ -179,7 +179,7 @@ class TLSLayer(layer.Layer): if isinstance(command, commands.SendData): yield from self.send(command) elif isinstance(command, commands.OpenConnection): - raise NotImplementedError() + raise NotImplementedError("Cannot open connection") else: yield command @@ -217,7 +217,7 @@ class TLSLayer(layer.Layer): self.recv_buffer[event.connection].extend(event.data) if event.connection == client and self.parse_client_hello(): - self._debug("SNI", self.client_hello.sni) + yield commands.Log(f"Client Hello: {self.client_hello}") client_tls_requires_server_connection = ( self.context.server.tls and diff --git a/mitmproxy/proxy2/server.py b/mitmproxy/proxy2/server.py index dac5508e2..173888f34 100644 --- a/mitmproxy/proxy2/server.py +++ b/mitmproxy/proxy2/server.py @@ -11,7 +11,7 @@ import asyncio import socket import typing -from mitmproxy import options +from mitmproxy import options as moptions from mitmproxy.proxy2 import events, commands from mitmproxy.proxy2.context import Client, Context, Connection from mitmproxy.proxy2.layers.modes import ReverseProxy @@ -40,14 +40,13 @@ class ConnectionHandler(metaclass=abc.ABCMeta): self.client: StreamIO(reader, writer) } - def _debug(self, *args): - print(*args) - async def handle_client(self): + self.log("clientconnect") + self.server_event(events.Start()) await self.handle_connection(self.client) - self._debug("client connection done, closing transports!") + self.log("clientdisconnect") if self.transports: await asyncio.wait([ @@ -55,13 +54,11 @@ class ConnectionHandler(metaclass=abc.ABCMeta): for x in self.transports ]) - self._debug("transports closed!") + # self._debug("transports closed!") async def close_connection(self, connection): io = self.transports.pop(connection, None) - if not io: - self._debug(f"Already closed: {connection}") - self._debug(f"Closing {connection}") + self.log(f"closing {connection}", "debug") try: await io.w.drain() io.w.write_eof() @@ -93,20 +90,25 @@ class ConnectionHandler(metaclass=abc.ABCMeta): except IOError as e: self.server_event(events.OpenConnectionReply(command, str(e))) else: + self.log("serverconnect") self.transports[command.connection] = StreamIO(reader, writer) command.connection.connected = True self.server_event(events.OpenConnectionReply(command, None)) await self.handle_connection(command.connection) + self.log("serverdisconnect") @abc.abstractmethod async def handle_hook(self, hook: commands.Hook) -> None: pass + def log(self, message: str, level: str = "info") -> None: + print(message) + def server_event(self, event: events.Event) -> None: - self._debug(">>", event) + self.log(f">> {event}", "debug") layer_commands = self.layer.handle_event(event) for command in layer_commands: - self._debug("<<", command) + self.log(f"<< {command}", "debug") if isinstance(command, commands.OpenConnection): asyncio.ensure_future( self.open_connection(command) @@ -121,8 +123,10 @@ class ConnectionHandler(metaclass=abc.ABCMeta): asyncio.ensure_future( self.handle_hook(command) ) + elif isinstance(command, commands.Log): + self.log(command.message, command.level) else: - raise RuntimeError(f"Unexpected event: {command}") + raise RuntimeError(f"Unexpected command: {command}") class SimpleConnectionHandler(ConnectionHandler): @@ -136,8 +140,11 @@ class SimpleConnectionHandler(ConnectionHandler): if __name__ == "__main__": loop = asyncio.get_event_loop() + opts = moptions.Options() + opts.mode = "reverse:example.com" + async def handle(reader, writer): - await SimpleConnectionHandler(reader, writer, options.Options()).handle_client() + await SimpleConnectionHandler(reader, writer, opts).handle_client() coro = asyncio.start_server(handle, '127.0.0.1', 8080, loop=loop) server = loop.run_until_complete(coro)