From ef582333ff432e11e696b95d7da456d8b6eae5cd Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 15 Mar 2017 12:47:03 +1300 Subject: [PATCH] Extract flow reading into addons This patch moves the final pieces of master functionality into addons. - Add a ReadFile addon to read from file - Add a separate ReadStdin addon to read from stdin, only used by mitmdump - Remove all methods that know about io and serialization from master.Master --- mitmproxy/addons/__init__.py | 2 + mitmproxy/addons/readfile.py | 50 +++++++++++++++++ mitmproxy/addons/readstdin.py | 34 +++++++++++ mitmproxy/master.py | 30 ---------- mitmproxy/tools/console/master.py | 20 ++----- mitmproxy/tools/dump.py | 17 +----- mitmproxy/tools/web/app.py | 3 +- mitmproxy/tools/web/master.py | 9 --- test/mitmproxy/addons/test_readfile.py | 62 +++++++++++++++++++++ test/mitmproxy/addons/test_readstdin.py | 59 ++++++++++++++++++++ test/mitmproxy/test_flow.py | 53 +++++------------- test/mitmproxy/tools/console/test_master.py | 6 +- test/mitmproxy/tools/test_dump.py | 13 ----- 13 files changed, 235 insertions(+), 123 deletions(-) create mode 100644 mitmproxy/addons/readfile.py create mode 100644 mitmproxy/addons/readstdin.py create mode 100644 test/mitmproxy/addons/test_readfile.py create mode 100644 test/mitmproxy/addons/test_readstdin.py diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index 7a45106c0..b4367d785 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -8,6 +8,7 @@ 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 @@ -37,5 +38,6 @@ 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 new file mode 100644 index 000000000..a4b924445 --- /dev/null +++ b/mitmproxy/addons/readfile.py @@ -0,0 +1,50 @@ +import os.path + +from mitmproxy import ctx +from mitmproxy import io +from mitmproxy import exceptions + + +class ReadFile: + """ + An addon that handles reading from file on startup. + """ + def __init__(self): + self.path = None + self.keepserving = False + + def load_flows_file(self, path: str) -> int: + path = os.path.expanduser(path) + cnt = 0 + try: + with open(path, "rb") as f: + freader = io.FlowReader(f) + for i in freader.stream(): + cnt += 1 + ctx.master.load_flow(i) + return cnt + except (IOError, exceptions.FlowReadException) as v: + if cnt: + ctx.log.warn( + "Flow file corrupted - loaded %i flows." % cnt, + ) + else: + ctx.log.error("Flow file corrupted.") + raise exceptions.FlowReadException(v) + + def configure(self, options, updated): + if "keepserving" in updated: + self.keepserving = options.keepserving + if "rfile" in updated and options.rfile: + self.path = options.rfile + + def running(self): + if self.path: + try: + self.load_flows_file(self.path) + except exceptions.FlowReadException as v: + raise exceptions.OptionsError(v) + finally: + self.path = None + if not self.keepserving: + ctx.master.shutdown() diff --git a/mitmproxy/addons/readstdin.py b/mitmproxy/addons/readstdin.py new file mode 100644 index 000000000..e45d25b8b --- /dev/null +++ b/mitmproxy/addons/readstdin.py @@ -0,0 +1,34 @@ +from mitmproxy import ctx +from mitmproxy import io +from mitmproxy import exceptions +import sys + + +class ReadStdin: + """ + An addon that reads from stdin if we're not attached to (someting like) + a tty. + """ + def __init__(self): + self.keepserving = False + + def configure(self, options, updated): + if "keepserving" in updated: + self.keepserving = options.keepserving + + def running(self, stdin = sys.stdin): + if not stdin.isatty(): + ctx.log.info("Reading from stdin") + try: + stdin.buffer.read(0) + except Exception as e: + ctx.log.warn("Cannot read from stdin: {}".format(e)) + return + freader = io.FlowReader(stdin.buffer) + try: + for i in freader.stream(): + ctx.master.load_flow(i) + except exceptions.FlowReadException as e: + ctx.log.error("Error reading from stdin: %s" % e) + if not self.keepserving: + ctx.master.shutdown() diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 79747a97d..69359de61 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -1,8 +1,6 @@ -import os import threading import contextlib import queue -import sys from mitmproxy import addonmanager from mitmproxy import options @@ -12,7 +10,6 @@ from mitmproxy import exceptions from mitmproxy import connections from mitmproxy import http from mitmproxy import log -from mitmproxy import io from mitmproxy.proxy.protocol import http_replay from mitmproxy.types import basethread import mitmproxy.net.http @@ -160,33 +157,6 @@ class Master: for e, o in eventsequence.iterate(f): getattr(self, e)(o) - def load_flows(self, fr: io.FlowReader) -> int: - """ - Load flows from a FlowReader object. - """ - cnt = 0 - for i in fr.stream(): - cnt += 1 - self.load_flow(i) - return cnt - - def load_flows_file(self, path: str) -> int: - path = os.path.expanduser(path) - try: - if path == "-": - try: - sys.stdin.buffer.read(0) - except Exception as e: - raise IOError("Cannot read from stdin: {}".format(e)) - freader = io.FlowReader(sys.stdin.buffer) - return self.load_flows(freader) - else: - with open(path, "rb") as f: - freader = io.FlowReader(f) - return self.load_flows(freader) - except IOError as v: - raise exceptions.FlowReadException(v.strerror) - def replay_request( self, f: http.HTTPFlow, diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index e75105cfc..b6339817c 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -256,19 +256,6 @@ class ConsoleMaster(master.Master): ) self.ab = statusbar.ActionBar() - if self.options.rfile: - ret = self.load_flows_path(self.options.rfile) - if ret and self.view.store_count(): - signals.add_log( - "File truncated or corrupted. " - "Loaded as many flows as possible.", - "error" - ) - elif ret and not self.view.store_count(): - self.shutdown() - print("Could not load file: {}".format(ret), file=sys.stderr) - sys.exit(1) - self.loop.set_alarm_in(0.01, self.ticker) self.loop.set_alarm_in( @@ -289,7 +276,10 @@ class ConsoleMaster(master.Master): print("Shutting down...", file=sys.stderr) finally: sys.stderr.flush() - self.shutdown() + super().shutdown() + + def shutdown(self): + raise urwid.ExitMainLoop def view_help(self, helpctx): signals.push_view_state.send( @@ -402,7 +392,7 @@ class ConsoleMaster(master.Master): def quit(self, a): if a != "n": - raise urwid.ExitMainLoop + self.shutdown() def clear_events(self): self.logbuffer[:] = [] diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py index 4bfe2dc49..be83fb1df 100644 --- a/mitmproxy/tools/dump.py +++ b/mitmproxy/tools/dump.py @@ -1,9 +1,8 @@ from mitmproxy import controller -from mitmproxy import exceptions from mitmproxy import addons from mitmproxy import options from mitmproxy import master -from mitmproxy.addons import dumper, termlog, termstatus +from mitmproxy.addons import dumper, termlog, termstatus, readstdin class DumpMaster(master.Master): @@ -22,21 +21,9 @@ class DumpMaster(master.Master): self.addons.add(*addons.default_addons()) if with_dumper: self.addons.add(dumper.Dumper()) - - if options.rfile: - try: - self.load_flows_file(options.rfile) - except exceptions.FlowReadException as v: - self.add_log("Flow file corrupted.", "error") - raise exceptions.OptionsError(v) + self.addons.add(readstdin.ReadStdin()) @controller.handler def log(self, e): if e.level == "error": self.has_errored = True - - def run(self): # pragma: no cover - if self.options.rfile and not self.options.keepserving: - self.addons.done() - return - super().run() diff --git a/mitmproxy/tools/web/app.py b/mitmproxy/tools/web/app.py index eddaa3e10..002513b9b 100644 --- a/mitmproxy/tools/web/app.py +++ b/mitmproxy/tools/web/app.py @@ -230,7 +230,8 @@ class DumpFlows(RequestHandler): def post(self): self.view.clear() bio = BytesIO(self.filecontents) - self.master.load_flows(io.FlowReader(bio)) + for i in io.FlowReader(bio).stream(): + self.master.load_flow(i) bio.close() diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index 8c7f579d8..e28bd0025 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -3,7 +3,6 @@ import webbrowser import tornado.httpserver import tornado.ioloop from mitmproxy import addons -from mitmproxy import exceptions from mitmproxy import log from mitmproxy import master from mitmproxy.addons import eventstore @@ -42,14 +41,6 @@ class WebMaster(master.Master): ) # This line is just for type hinting self.options = self.options # type: Options - if options.rfile: - try: - self.load_flows_file(options.rfile) - except exceptions.FlowReadException as v: - self.add_log( - "Could not read flow file: %s" % v, - "error" - ) def _sig_view_add(self, view, flow): app.ClientConnection.broadcast( diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py new file mode 100644 index 000000000..c0cf97ae3 --- /dev/null +++ b/test/mitmproxy/addons/test_readfile.py @@ -0,0 +1,62 @@ +from mitmproxy.addons import readfile +from mitmproxy.test import taddons +from mitmproxy.test import tflow +from mitmproxy import io +from mitmproxy import exceptions +from unittest import mock + +import pytest + + +def write_data(path, corrupt=False): + with open(path, "wb") as tf: + w = io.FlowWriter(tf) + for i in range(3): + f = tflow.tflow(resp=True) + w.add(f) + for i in range(3): + f = tflow.tflow(err=True) + w.add(f) + f = tflow.ttcpflow() + w.add(f) + f = tflow.ttcpflow(err=True) + w.add(f) + if corrupt: + tf.write(b"flibble") + + +@mock.patch('mitmproxy.master.Master.load_flow') +def test_configure(mck, tmpdir): + + rf = readfile.ReadFile() + with taddons.context() as tctx: + tf = str(tmpdir.join("tfile")) + write_data(tf) + tctx.configure(rf, rfile=str(tf), keepserving=False) + assert not mck.called + rf.running() + assert mck.called + + write_data(tf, corrupt=True) + tctx.configure(rf, rfile=str(tf), keepserving=False) + with pytest.raises(exceptions.OptionsError): + rf.running() + + +@mock.patch('mitmproxy.master.Master.load_flow') +def test_corruption(mck, tmpdir): + + rf = readfile.ReadFile() + with taddons.context() as tctx: + with pytest.raises(exceptions.FlowReadException): + rf.load_flows_file("nonexistent") + assert not mck.called + assert len(tctx.master.event_log) == 1 + + tfc = str(tmpdir.join("tfile")) + write_data(tfc, corrupt=True) + + with pytest.raises(exceptions.FlowReadException): + rf.load_flows_file(tfc) + assert mck.called + assert len(tctx.master.event_log) == 2 diff --git a/test/mitmproxy/addons/test_readstdin.py b/test/mitmproxy/addons/test_readstdin.py new file mode 100644 index 000000000..bbef81fce --- /dev/null +++ b/test/mitmproxy/addons/test_readstdin.py @@ -0,0 +1,59 @@ + +import io +from mitmproxy.addons import readstdin +from mitmproxy.test import taddons +from mitmproxy.test import tflow +import mitmproxy.io +from unittest import mock + + +def gen_data(corrupt=False): + tf = io.BytesIO() + w = mitmproxy.io.FlowWriter(tf) + for i in range(3): + f = tflow.tflow(resp=True) + w.add(f) + for i in range(3): + f = tflow.tflow(err=True) + w.add(f) + f = tflow.ttcpflow() + w.add(f) + f = tflow.ttcpflow(err=True) + w.add(f) + if corrupt: + tf.write(b"flibble") + tf.seek(0) + return tf + + +def test_configure(tmpdir): + rf = readstdin.ReadStdin() + with taddons.context() as tctx: + tctx.configure(rf, keepserving=False) + + +class mStdin: + def __init__(self, d): + self.buffer = d + + def isatty(self): + return False + + +@mock.patch('mitmproxy.master.Master.load_flow') +def test_read(m, tmpdir): + rf = readstdin.ReadStdin() + with taddons.context() as tctx: + assert not m.called + rf.running(stdin=mStdin(gen_data())) + assert m.called + + rf.running(stdin=mStdin(None)) + assert tctx.master.event_log + tctx.master.clear() + + m.reset_mock() + assert not m.called + rf.running(stdin=mStdin(gen_data(corrupt=True))) + assert m.called + assert tctx.master.event_log diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index f4d32cbb9..1fb33bb22 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -3,12 +3,13 @@ import pytest from mitmproxy.test import tflow import mitmproxy.io -from mitmproxy import flowfilter, options +from mitmproxy import flowfilter +from mitmproxy import options +from mitmproxy.proxy import config from mitmproxy.contrib import tnetstring from mitmproxy.exceptions import FlowReadException from mitmproxy import flow from mitmproxy import http -from mitmproxy.proxy import ProxyConfig from mitmproxy.proxy.server import DummyServer from mitmproxy import master from . import tservers @@ -16,23 +17,6 @@ from . import tservers class TestSerialize: - def _treader(self): - sio = io.BytesIO() - w = mitmproxy.io.FlowWriter(sio) - for i in range(3): - f = tflow.tflow(resp=True) - w.add(f) - for i in range(3): - f = tflow.tflow(err=True) - w.add(f) - f = tflow.ttcpflow() - w.add(f) - f = tflow.ttcpflow(err=True) - w.add(f) - - sio.seek(0) - return mitmproxy.io.FlowReader(sio) - def test_roundtrip(self): sio = io.BytesIO() f = tflow.tflow() @@ -51,26 +35,6 @@ class TestSerialize: assert f2.request == f.request assert f2.marked - def test_load_flows(self): - r = self._treader() - s = tservers.TestState() - fm = master.Master(None, DummyServer()) - fm.addons.add(s) - fm.load_flows(r) - assert len(s.flows) == 6 - - def test_load_flows_reverse(self): - r = self._treader() - s = tservers.TestState() - opts = options.Options( - mode="reverse:https://use-this-domain" - ) - conf = ProxyConfig(opts) - fm = master.Master(opts, DummyServer(conf)) - fm.addons.add(s) - fm.load_flows(r) - assert s.flows[0].request.host == "use-this-domain" - def test_filter(self): sio = io.BytesIO() flt = flowfilter.parse("~c 200") @@ -122,6 +86,17 @@ class TestSerialize: class TestFlowMaster: + def test_load_flow_reverse(self): + s = tservers.TestState() + opts = options.Options( + mode="reverse:https://use-this-domain" + ) + conf = config.ProxyConfig(opts) + fm = master.Master(opts, DummyServer(conf)) + fm.addons.add(s) + f = tflow.tflow(resp=True) + fm.load_flow(f) + assert s.flows[0].request.host == "use-this-domain" def test_replay(self): fm = master.Master(None, DummyServer()) diff --git a/test/mitmproxy/tools/console/test_master.py b/test/mitmproxy/tools/console/test_master.py index 6c716ad1b..459084500 100644 --- a/test/mitmproxy/tools/console/test_master.py +++ b/test/mitmproxy/tools/console/test_master.py @@ -5,6 +5,7 @@ from mitmproxy import proxy from mitmproxy import options from mitmproxy.tools.console import common from ... import tservers +import urwid def test_format_keyvals(): @@ -35,7 +36,10 @@ class TestMaster(tservers.MasterTest): def test_basic(self): m = self.mkmaster() for i in (1, 2, 3): - self.dummy_cycle(m, 1, b"") + try: + self.dummy_cycle(m, 1, b"") + except urwid.ExitMainLoop: + pass assert len(m.view) == i def test_run_script_once(self): diff --git a/test/mitmproxy/tools/test_dump.py b/test/mitmproxy/tools/test_dump.py index a15bf5830..624bf08f4 100644 --- a/test/mitmproxy/tools/test_dump.py +++ b/test/mitmproxy/tools/test_dump.py @@ -2,7 +2,6 @@ import pytest from unittest import mock from mitmproxy import proxy -from mitmproxy import exceptions from mitmproxy import log from mitmproxy import controller from mitmproxy import options @@ -17,18 +16,6 @@ class TestDumpMaster(tservers.MasterTest): m = dump.DumpMaster(o, proxy.DummyServer(), with_termlog=False, with_dumper=False) return m - def test_read(self, tmpdir): - p = str(tmpdir.join("read")) - self.flowfile(p) - self.dummy_cycle( - self.mkmaster(None, rfile=p), - 1, b"", - ) - with pytest.raises(exceptions.OptionsError): - self.mkmaster(None, rfile="/nonexistent") - with pytest.raises(exceptions.OptionsError): - self.mkmaster(None, rfile="test_dump.py") - def test_has_error(self): m = self.mkmaster(None) ent = log.LogEntry("foo", "error")