From 7fe09f4d26e604db9d3e7f9868b3f0da0b3a6406 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Tue, 4 Oct 2016 13:00:15 +1100 Subject: [PATCH] Consolidate logging - Make log event handler work - Move terminal event log into an addon - Clean various log related errors and duplications --- mitmproxy/builtins/termlog.py | 23 ++++++++++++++++++++ mitmproxy/controller.py | 20 ++++++++++++----- mitmproxy/dump.py | 19 +++++----------- mitmproxy/flow/master.py | 2 +- mitmproxy/proxy/__init__.py | 4 ++-- mitmproxy/proxy/root_context.py | 9 ++------ mitmproxy/proxy/server.py | 3 ++- test/mitmproxy/builtins/test_termlog.py | 17 +++++++++++++++ test/mitmproxy/data/addonscripts/recorder.py | 4 ++-- test/mitmproxy/data/scripts/all.py | 16 +++++++------- test/mitmproxy/mastertest.py | 5 +---- test/mitmproxy/tservers.py | 5 +++-- 12 files changed, 81 insertions(+), 46 deletions(-) create mode 100644 mitmproxy/builtins/termlog.py create mode 100644 test/mitmproxy/builtins/test_termlog.py diff --git a/mitmproxy/builtins/termlog.py b/mitmproxy/builtins/termlog.py new file mode 100644 index 000000000..c8259f0c1 --- /dev/null +++ b/mitmproxy/builtins/termlog.py @@ -0,0 +1,23 @@ +from __future__ import absolute_import, print_function, division + +import click + +from mitmproxy import utils + + +class TermLog: + def __init__(self): + self.options = None + + def configure(self, options, updated): + self.options = options + + def log(self, e): + if self.options.verbosity >= utils.log_tier(e.level): + click.secho( + e.msg, + file=self.options.tfile, + fg=dict(error="red", warn="yellow").get(e.level), + dim=(e.level == "debug"), + err=(e.level == "error") + ) diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 51a9b1460..1a8801e1f 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -28,21 +28,26 @@ Events = frozenset([ "requestheaders", "response", "responseheaders", + "error", "websocket_handshake", "next_layer", - "error", - "log", - - "start", "configure", "done", + "log", + "start", "tick", ]) +class LogEntry(object): + def __init__(self, msg, level): + self.msg = msg + self.level = level + + class Log(object): def __init__(self, master): self.master = master @@ -90,11 +95,16 @@ class Master(object): mitmproxy_ctx.master = None mitmproxy_ctx.log = None + def tell(self, mtype, m): + m.reply = DummyReply() + self.event_queue.put((mtype, m)) + def add_log(self, e, level): """ level: debug, info, warn, error """ - pass + with self.handlecontext(): + self.addons("log", LogEntry(e, level)) def add_server(self, server): # We give a Channel to the server which can be used to communicate with the master diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index 6bed8fcc7..0e942eacf 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -3,15 +3,12 @@ from __future__ import absolute_import, print_function, division from typing import Optional # noqa import typing # noqa -import click - from mitmproxy import controller from mitmproxy import exceptions from mitmproxy import flow from mitmproxy import builtins -from mitmproxy import utils from mitmproxy import options -from mitmproxy.builtins import dumper +from mitmproxy.builtins import dumper, termlog from netlib import tcp @@ -42,6 +39,7 @@ class DumpMaster(flow.FlowMaster): self.has_errored = False self.addons.add(*builtins.default_addons()) self.addons.add(dumper.Dumper()) + self.addons.add(termlog.TermLog()) # This line is just for type hinting self.options = self.options # type: Options self.set_stream_large_bodies(options.stream_large_bodies) @@ -79,17 +77,10 @@ class DumpMaster(flow.FlowMaster): except exceptions.FlowReadException as e: raise DumpError(str(e)) - def add_log(self, e, level="info"): - if level == "error": + @controller.handler + def log(self, e): + if e.level == "error": self.has_errored = True - if self.options.verbosity >= utils.log_tier(level): - click.secho( - e, - file=self.options.tfile, - fg=dict(error="red", warn="yellow").get(level), - dim=(level == "debug"), - err=(level == "error") - ) @controller.handler def request(self, f): diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index 7149fec91..a2b225b86 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -190,7 +190,7 @@ class FlowMaster(controller.Master): @controller.handler def log(self, l): - self.add_log(l.msg, l.level) + pass @controller.handler def clientconnect(self, root_layer): diff --git a/mitmproxy/proxy/__init__.py b/mitmproxy/proxy/__init__.py index ada9fa12e..e3d249fa0 100644 --- a/mitmproxy/proxy/__init__.py +++ b/mitmproxy/proxy/__init__.py @@ -1,11 +1,11 @@ from __future__ import absolute_import, print_function, division from .config import ProxyConfig -from .root_context import RootContext, Log +from .root_context import RootContext from .server import ProxyServer, DummyServer __all__ = [ "ProxyServer", "DummyServer", "ProxyConfig", - "RootContext", "Log", + "RootContext" ] diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py index 9827ef744..cfa04bdf7 100644 --- a/mitmproxy/proxy/root_context.py +++ b/mitmproxy/proxy/root_context.py @@ -5,6 +5,7 @@ import sys import six import netlib.exceptions +from mitmproxy import controller from mitmproxy import exceptions from mitmproxy import protocol from mitmproxy.proxy import modes @@ -118,7 +119,7 @@ class RootContext(object): for i in subs: full_msg.append(" -> " + i) full_msg = "\n".join(full_msg) - self.channel.tell("log", Log(full_msg, level)) + self.channel.tell("log", controller.LogEntry(full_msg, level)) @property def layers(self): @@ -126,9 +127,3 @@ class RootContext(object): def __repr__(self): return "RootContext" - - -class Log(object): - def __init__(self, msg, level="info"): - self.msg = msg - self.level = level diff --git a/mitmproxy/proxy/server.py b/mitmproxy/proxy/server.py index c5fd5f9e7..182f46024 100644 --- a/mitmproxy/proxy/server.py +++ b/mitmproxy/proxy/server.py @@ -9,6 +9,7 @@ import six import netlib.exceptions from mitmproxy import exceptions from mitmproxy import models +from mitmproxy import controller from mitmproxy.proxy import modes from mitmproxy.proxy import root_context from netlib import tcp @@ -155,4 +156,4 @@ class ConnectionHandler(object): def log(self, msg, level): msg = "{}: {}".format(repr(self.client_conn.address), msg) - self.channel.tell("log", root_context.Log(msg, level)) + self.channel.tell("log", controller.LogEntry(msg, level)) diff --git a/test/mitmproxy/builtins/test_termlog.py b/test/mitmproxy/builtins/test_termlog.py new file mode 100644 index 000000000..a15104426 --- /dev/null +++ b/test/mitmproxy/builtins/test_termlog.py @@ -0,0 +1,17 @@ +from .. import mastertest +from six.moves import cStringIO as StringIO + +from mitmproxy.builtins import termlog +from mitmproxy import controller +from mitmproxy import dump + + +class TestTermLog(mastertest.MasterTest): + def test_simple(self): + t = termlog.TermLog() + sio = StringIO() + t.configure(dump.Options(tfile = sio, verbosity = 2), set([])) + t.log(controller.LogEntry("one", "info")) + assert "one" in sio.getvalue() + t.log(controller.LogEntry("two", "debug")) + assert "two" not in sio.getvalue() diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py index 890e6f4e5..3d7a51bf9 100644 --- a/test/mitmproxy/data/addonscripts/recorder.py +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -15,8 +15,8 @@ class CallLogger: lg = (self.name, attr, args, kwargs) if attr != "log": ctx.log.info(str(lg)) - self.call_log.append(lg) - ctx.log.debug("%s %s" % (self.name, attr)) + self.call_log.append(lg) + ctx.log.debug("%s %s" % (self.name, attr)) return prox raise AttributeError diff --git a/test/mitmproxy/data/scripts/all.py b/test/mitmproxy/data/scripts/all.py index bf8e93ecf..95edfe977 100644 --- a/test/mitmproxy/data/scripts/all.py +++ b/test/mitmproxy/data/scripts/all.py @@ -1,37 +1,37 @@ import mitmproxy -log = [] +record = [] def clientconnect(cc): mitmproxy.ctx.log("XCLIENTCONNECT") - log.append("clientconnect") + record.append("clientconnect") def serverconnect(cc): mitmproxy.ctx.log("XSERVERCONNECT") - log.append("serverconnect") + record.append("serverconnect") def request(f): mitmproxy.ctx.log("XREQUEST") - log.append("request") + record.append("request") def response(f): mitmproxy.ctx.log("XRESPONSE") - log.append("response") + record.append("response") def responseheaders(f): mitmproxy.ctx.log("XRESPONSEHEADERS") - log.append("responseheaders") + record.append("responseheaders") def clientdisconnect(cc): mitmproxy.ctx.log("XCLIENTDISCONNECT") - log.append("clientdisconnect") + record.append("clientdisconnect") def error(cc): mitmproxy.ctx.log("XERROR") - log.append("error") + record.append("error") diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py index a14fe02a5..5605b1a6e 100644 --- a/test/mitmproxy/mastertest.py +++ b/test/mitmproxy/mastertest.py @@ -4,7 +4,7 @@ from . import tutils import netlib.tutils from mitmproxy.flow import master -from mitmproxy import flow, proxy, models, controller, options +from mitmproxy import flow, proxy, models, options class TestMaster: @@ -15,9 +15,6 @@ class MasterTest: def cycle(self, master, content): f = tutils.tflow(req=netlib.tutils.treq(content=content)) - l = proxy.Log("connect") - l.reply = controller.DummyReply() - master.log(l) master.clientconnect(f.client_conn) master.serverconnect(f.server_conn) master.request(f) diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 119c2caea..d3806a997 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -44,8 +44,9 @@ class TestMaster(flow.FlowMaster): def clear_log(self): self.tlog = [] - def add_log(self, message, level=None): - self.tlog.append(message) + @controller.handler + def log(self, e): + self.tlog.append(e.msg) class ProxyThread(threading.Thread):