diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py index 42930a7e1..adb1f7ef5 100644 --- a/mitmproxy/tools/dump.py +++ b/mitmproxy/tools/dump.py @@ -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 + ) diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 35567b620..b321e8f80 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -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 diff --git a/test/mitmproxy/tools/test_dump.py b/test/mitmproxy/tools/test_dump.py index 624bf08f4..810fcaaad 100644 --- a/test/mitmproxy/tools/test_dump.py +++ b/test/mitmproxy/tools/test_dump.py @@ -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):