mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-30 03:14:22 +00:00
Move dump handlers to addon
This commit is contained in:
parent
b275257553
commit
b6abe96202
@ -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
|
||||
)
|
||||
|
@ -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
|
||||
|
@ -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):
|
||||
|
Loading…
Reference in New Issue
Block a user