Move dump handlers to addon

This commit is contained in:
Aldo Cortesi 2017-03-16 11:18:23 +13:00 committed by Aldo Cortesi
parent b275257553
commit b6abe96202
3 changed files with 21 additions and 13 deletions

View File

@ -1,10 +1,19 @@
from mitmproxy import controller
from mitmproxy import ctx
from mitmproxy import addons
from mitmproxy import options
from mitmproxy import master
from mitmproxy.addons import dumper, termlog, termstatus, readstdin, keepserving
class ErrorCheck:
def __init__(self):
self.has_errored = False
def log(self, e):
if e.level == "error":
self.has_errored = True
class DumpMaster(master.Master):
def __init__(
@ -15,15 +24,14 @@ class DumpMaster(master.Master):
with_dumper=True,
) -> None:
master.Master.__init__(self, options, server)
self.has_errored = False
self.errorcheck = ErrorCheck()
if with_termlog:
self.addons.add(termlog.TermLog(), termstatus.TermStatus())
self.addons.add(*addons.default_addons())
if with_dumper:
self.addons.add(dumper.Dumper())
self.addons.add(readstdin.ReadStdin(), keepserving.KeepServing())
@controller.handler
def log(self, e):
if e.level == "error":
self.has_errored = True
self.addons.add(
readstdin.ReadStdin(),
keepserving.KeepServing(),
self.errorcheck
)

View File

@ -91,9 +91,7 @@ def run(MasterKlass, args): # pragma: no cover
sys.exit(1)
except (KeyboardInterrupt, RuntimeError):
pass
if master is None or getattr(master, "has_errored", None):
print("%s: errors occurred during run" % sys.argv[0], file=sys.stderr)
sys.exit(1)
return master
def mitmproxy(args=None): # pragma: no cover
@ -109,7 +107,9 @@ def mitmproxy(args=None): # pragma: no cover
def mitmdump(args=None): # pragma: no cover
from mitmproxy.tools import dump
run(dump.DumpMaster, args)
m = run(dump.DumpMaster, args)
if m and m.errorcheck.has_errored:
sys.exit(1)
def mitmweb(args=None): # pragma: no cover

View File

@ -21,7 +21,7 @@ class TestDumpMaster(tservers.MasterTest):
ent = log.LogEntry("foo", "error")
ent.reply = controller.DummyReply()
m.log(ent)
assert m.has_errored
assert m.errorcheck.has_errored
@pytest.mark.parametrize("termlog", [False, True])
def test_addons_termlog(self, termlog):