[sans-io] improve logging

This commit is contained in:
Maximilian Hils 2017-10-26 18:32:03 +02:00
parent b5a3343d03
commit edaf3219fc
7 changed files with 47 additions and 43 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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