Move dump handlers to addon

3 files changed, 21 insertions, 13 deletions

diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py
index 42930a7e..adb1f7ef 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 35567b62..b321e8f8 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 624bf08f..810fcaaa 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):