From ca2827886a97de88f4ab2937e71588fc9320ba4d Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Wed, 26 Apr 2017 11:12:14 +0200 Subject: [PATCH] separate reading from stdin into its own addon --- mitmproxy/addons/__init__.py | 2 - mitmproxy/addons/readfile.py | 32 +++++--- mitmproxy/tools/console/master.py | 10 ++- mitmproxy/tools/dump.py | 3 +- mitmproxy/tools/web/master.py | 2 + test/mitmproxy/addons/test_readfile.py | 107 ++++++++++++++----------- 6 files changed, 89 insertions(+), 67 deletions(-) diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index b4367d785..7a45106c0 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -8,7 +8,6 @@ from mitmproxy.addons import disable_h2c from mitmproxy.addons import onboarding from mitmproxy.addons import proxyauth from mitmproxy.addons import replace -from mitmproxy.addons import readfile from mitmproxy.addons import script from mitmproxy.addons import serverplayback from mitmproxy.addons import setheaders @@ -38,6 +37,5 @@ def default_addons(): stickycookie.StickyCookie(), streambodies.StreamBodies(), streamfile.StreamFile(), - readfile.ReadFile(), upstream_auth.UpstreamAuth(), ] diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py index e83c99934..05b6c3099 100644 --- a/mitmproxy/addons/readfile.py +++ b/mitmproxy/addons/readfile.py @@ -1,17 +1,17 @@ import os.path +import sys import typing -import sys - from mitmproxy import ctx -from mitmproxy import io from mitmproxy import exceptions +from mitmproxy import io class ReadFile: """ An addon that handles reading from file on startup. """ + def load_flows(self, fo: typing.IO[bytes]) -> int: cnt = 0 freader = io.FlowReader(fo) @@ -29,16 +29,13 @@ class ReadFile: return cnt def load_flows_from_path(self, path: str) -> int: - if path == "-": - return self.load_flows(sys.stdin.buffer) - else: - path = os.path.expanduser(path) - try: - with open(path, "rb") as f: - return self.load_flows(f) - except IOError as e: - ctx.log.error("Cannot load flows: {}".format(e)) - raise exceptions.FlowReadException(str(e)) from e + path = os.path.expanduser(path) + try: + with open(path, "rb") as f: + return self.load_flows(f) + except IOError as e: + ctx.log.error("Cannot load flows: {}".format(e)) + raise exceptions.FlowReadException(str(e)) from e def running(self): if ctx.options.rfile: @@ -48,3 +45,12 @@ class ReadFile: raise exceptions.OptionsError(e) from e finally: ctx.master.addons.trigger("processing_complete") + + +class ReadFileStdin(ReadFile): + """Support the special case of "-" for reading from stdin""" + def load_flows_from_path(self, path: str) -> int: + if path == "-": + return self.load_flows(sys.stdin.buffer) + else: + return super().load_flows_from_path(path) diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index c1d584acc..e7a2c6ae4 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -17,8 +17,9 @@ from mitmproxy import exceptions from mitmproxy import master from mitmproxy import io from mitmproxy import log -from mitmproxy.addons import view from mitmproxy.addons import intercept +from mitmproxy.addons import readfile +from mitmproxy.addons import view from mitmproxy.tools.console import flowlist from mitmproxy.tools.console import flowview from mitmproxy.tools.console import grideditor @@ -91,7 +92,12 @@ class ConsoleMaster(master.Master): signals.sig_add_log.connect(self.sig_add_log) self.addons.add(Logger()) self.addons.add(*addons.default_addons()) - self.addons.add(intercept.Intercept(), self.view, UnsupportedLog()) + self.addons.add( + intercept.Intercept(), + self.view, + UnsupportedLog(), + readfile.ReadFile(), + ) def sigint_handler(*args, **kwargs): self.prompt_for_exit() diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py index a4c9998bd..4d0ccf4bd 100644 --- a/mitmproxy/tools/dump.py +++ b/mitmproxy/tools/dump.py @@ -1,7 +1,7 @@ from mitmproxy import addons from mitmproxy import options from mitmproxy import master -from mitmproxy.addons import dumper, termlog, termstatus, keepserving +from mitmproxy.addons import dumper, termlog, termstatus, keepserving, readfile class ErrorCheck: @@ -31,5 +31,6 @@ class DumpMaster(master.Master): self.addons.add(dumper.Dumper()) self.addons.add( keepserving.KeepServing(), + readfile.ReadFileStdin(), self.errorcheck ) diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index 0db5a09f2..c09fe0a20 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -7,6 +7,7 @@ from mitmproxy import log from mitmproxy import master from mitmproxy.addons import eventstore from mitmproxy.addons import intercept +from mitmproxy.addons import readfile from mitmproxy.addons import termlog from mitmproxy.addons import view from mitmproxy.addons import termstatus @@ -32,6 +33,7 @@ class WebMaster(master.Master): self.addons.add(*addons.default_addons()) self.addons.add( intercept.Intercept(), + readfile.ReadFile(), self.view, self.events, ) diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py index 689d9779e..813aa10e5 100644 --- a/test/mitmproxy/addons/test_readfile.py +++ b/test/mitmproxy/addons/test_readfile.py @@ -37,59 +37,68 @@ def corrupt_data(): return f -@mock.patch('mitmproxy.master.Master.load_flow') -def test_configure(mck, tmpdir, data, corrupt_data): - rf = readfile.ReadFile() - with taddons.context() as tctx: - tf = tmpdir.join("tfile") +class TestReadFile: + @mock.patch('mitmproxy.master.Master.load_flow') + def test_configure(self, mck, tmpdir, data, corrupt_data): + rf = readfile.ReadFile() + with taddons.context() as tctx: + tf = tmpdir.join("tfile") - tf.write(data.getvalue()) - tctx.configure(rf, rfile=str(tf)) - assert not mck.called - rf.running() - assert mck.called - - tf.write(corrupt_data.getvalue()) - tctx.configure(rf, rfile=str(tf)) - with pytest.raises(exceptions.OptionsError): + tf.write(data.getvalue()) + tctx.configure(rf, rfile=str(tf)) + assert not mck.called rf.running() + assert mck.called + + tf.write(corrupt_data.getvalue()) + tctx.configure(rf, rfile=str(tf)) + with pytest.raises(exceptions.OptionsError): + rf.running() + + @mock.patch('mitmproxy.master.Master.load_flow') + def test_corrupt(self, mck, corrupt_data): + rf = readfile.ReadFile() + with taddons.context() as tctx: + with pytest.raises(exceptions.FlowReadException): + rf.load_flows(io.BytesIO(b"qibble")) + assert not mck.called + assert len(tctx.master.logs) == 1 + + with pytest.raises(exceptions.FlowReadException): + rf.load_flows(corrupt_data) + assert mck.called + assert len(tctx.master.logs) == 2 + + def test_nonexisting_file(self): + rf = readfile.ReadFile() + with taddons.context() as tctx: + with pytest.raises(exceptions.FlowReadException): + rf.load_flows_from_path("nonexistent") + assert len(tctx.master.logs) == 1 -@mock.patch('mitmproxy.master.Master.load_flow') -@mock.patch('sys.stdin') -def test_configure_stdin(stdin, load_flow, data, corrupt_data): - rf = readfile.ReadFile() - with taddons.context() as tctx: - stdin.buffer = data - tctx.configure(rf, rfile="-") - assert not load_flow.called - rf.running() - assert load_flow.called - - stdin.buffer = corrupt_data - tctx.configure(rf, rfile="-") - with pytest.raises(exceptions.OptionsError): +class TestReadFileStdin: + @mock.patch('mitmproxy.master.Master.load_flow') + @mock.patch('sys.stdin') + def test_stdin(self, stdin, load_flow, data, corrupt_data): + rf = readfile.ReadFileStdin() + with taddons.context() as tctx: + stdin.buffer = data + tctx.configure(rf, rfile="-") + assert not load_flow.called rf.running() + assert load_flow.called + stdin.buffer = corrupt_data + tctx.configure(rf, rfile="-") + with pytest.raises(exceptions.OptionsError): + rf.running() -@mock.patch('mitmproxy.master.Master.load_flow') -def test_corrupt(mck, corrupt_data): - rf = readfile.ReadFile() - with taddons.context() as tctx: - with pytest.raises(exceptions.FlowReadException): - rf.load_flows(io.BytesIO(b"qibble")) - assert not mck.called - assert len(tctx.master.logs) == 1 - - with pytest.raises(exceptions.FlowReadException): - rf.load_flows(corrupt_data) - assert mck.called - assert len(tctx.master.logs) == 2 - - -def test_nonexisting_file(): - rf = readfile.ReadFile() - with taddons.context() as tctx: - with pytest.raises(exceptions.FlowReadException): - rf.load_flows_from_path("nonexistent") - assert len(tctx.master.logs) == 1 + @mock.patch('mitmproxy.master.Master.load_flow') + def test_normal(self, load_flow, tmpdir, data): + rf = readfile.ReadFileStdin() + with taddons.context(): + tfile = tmpdir.join("tfile") + tfile.write(data.getvalue()) + rf.load_flows_from_path(str(tfile)) + assert load_flow.called