diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index c5d405253..d2b50c35e 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -1,7 +1,7 @@ from mitmproxy.addons import anticache from mitmproxy.addons import anticomp from mitmproxy.addons import clientplayback -from mitmproxy.addons import filestreamer +from mitmproxy.addons import streamfile from mitmproxy.addons import onboarding from mitmproxy.addons import replace from mitmproxy.addons import script @@ -20,7 +20,7 @@ def default_addons(): stickyauth.StickyAuth(), stickycookie.StickyCookie(), script.ScriptLoader(), - filestreamer.FileStreamer(), + streamfile.StreamFile(), streambodies.StreamBodies(), replace.Replace(), setheaders.SetHeaders(), diff --git a/mitmproxy/addons/onboardingapp/app.py b/mitmproxy/addons/onboardingapp/app.py index 50b52214a..d418952c7 100644 --- a/mitmproxy/addons/onboardingapp/app.py +++ b/mitmproxy/addons/onboardingapp/app.py @@ -6,7 +6,6 @@ import tornado.wsgi from mitmproxy.utils import data from mitmproxy.proxy import config -from mitmproxy.addons import wsgiapp loader = tornado.template.Loader(data.pkg_data.path("addons/onboardingapp/templates")) @@ -92,18 +91,3 @@ application = tornado.web.Application( ], # debug=True ) - - -class Onboarding(wsgiapp.WSGIApp): - def __init__(self): - super().__init__(Adapter(application), None, None) - self.enabled = False - - def configure(self, options, updated): - self.host = options.app_host - self.port = options.app_port - self.enabled = options.app - - def request(self, f): - if self.enabled: - super().request(f) diff --git a/mitmproxy/addons/filestreamer.py b/mitmproxy/addons/streamfile.py similarity index 72% rename from mitmproxy/addons/filestreamer.py rename to mitmproxy/addons/streamfile.py index 031b44ab6..377f277de 100644 --- a/mitmproxy/addons/filestreamer.py +++ b/mitmproxy/addons/streamfile.py @@ -5,9 +5,10 @@ from mitmproxy import flowfilter from mitmproxy import io -class FileStreamer: +class StreamFile: def __init__(self): self.stream = None + self.filt = None self.active_flows = set() # type: Set[flow.Flow] def start_stream_to_path(self, path, mode, flt): @@ -15,29 +16,28 @@ class FileStreamer: try: f = open(path, mode) except IOError as v: - return str(v) + raise exceptions.OptionsError(str(v)) self.stream = io.FilteredFlowWriter(f, flt) self.active_flows = set() def configure(self, options, updated): # We're already streaming - stop the previous stream and restart - if self.stream: - self.done() - - if options.outfile: - flt = None + if "filtstr" in updated: if options.get("filtstr"): - flt = flowfilter.parse(options.filtstr) - if not flt: + self.filt = flowfilter.parse(options.filtstr) + if not self.filt: raise exceptions.OptionsError( "Invalid filter specification: %s" % options.filtstr ) - path, mode = options.outfile - if mode not in ("wb", "ab"): - raise exceptions.OptionsError("Invalid mode.") - err = self.start_stream_to_path(path, mode, flt) - if err: - raise exceptions.OptionsError(err) + if "streamfile" in updated: + if self.stream: + self.done() + if options.streamfile: + if options.streamfile_append: + mode = "ab" + else: + mode = "wb" + self.start_stream_to_path(options.streamfile, mode, self.filt) def tcp_start(self, flow): if self.stream: diff --git a/mitmproxy/options.py b/mitmproxy/options.py index 1db9f0f08..497914006 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -48,7 +48,8 @@ class Options(optmanager.OptManager): stream_large_bodies: Optional[int] = None, verbosity: int = 2, default_contentview: str = "auto", - outfile: Optional[Tuple[str, str]] = None, + streamfile: Optional[str] = None, + streamfile_append: bool = False, server_replay_ignore_content: bool = False, server_replay_ignore_params: Sequence[str] = (), server_replay_ignore_payload_params: Sequence[str] = (), @@ -108,7 +109,8 @@ class Options(optmanager.OptManager): self.stream_large_bodies = stream_large_bodies self.verbosity = verbosity self.default_contentview = default_contentview - self.outfile = outfile + self.streamfile = streamfile + self.streamfile_append = streamfile_append self.server_replay_ignore_content = server_replay_ignore_content self.server_replay_ignore_params = server_replay_ignore_params self.server_replay_ignore_payload_params = server_replay_ignore_payload_params diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 3cba6762e..7804b90db 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -1,7 +1,10 @@ +import contextlib + import mitmproxy.master import mitmproxy.options from mitmproxy import proxy from mitmproxy import events +from mitmproxy import exceptions class RecordingMaster(mitmproxy.master.Master): @@ -36,6 +39,15 @@ class context: self.wrapped = None return False + @contextlib.contextmanager + def _rollback(self, opts, updates): + old = opts._opts.copy() + try: + yield + except exceptions.OptionsError as e: + opts.__dict__["_opts"] = old + raise + def cycle(self, addon, f): """ Cycles the flow through the events for the flow. Stops if a reply @@ -55,6 +67,6 @@ class context: Options object with the given keyword arguments, then calls the configure method on the addon with the updated value. """ - for k, v in kwargs.items(): - setattr(self.options, k, v) - addon.configure(self.options, kwargs.keys()) + with self._rollback(self.options, kwargs): + self.options.update(**kwargs) + addon.configure(self.options, kwargs.keys()) diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index e4b29d0fa..debe6db9f 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -140,8 +140,8 @@ def get_common_options(args): raise exceptions.OptionsError(e) setheaders.append(p) - if args.outfile and args.outfile[0] == args.rfile: - if args.outfile[1] == "wb": + if args.streamfile and args.streamfile[0] == args.rfile: + if args.streamfile[1] == "wb": raise exceptions.OptionsError( "Cannot use '{}' for both reading and writing flows. " "Are you looking for --afile?".format(args.rfile) @@ -228,7 +228,8 @@ def get_common_options(args): stickyauth=stickyauth, stream_large_bodies=stream_large_bodies, showhost=args.showhost, - outfile=args.outfile, + streamfile=args.streamfile[0] if args.streamfile else None, + streamfile_append=True if args.streamfile and args.streamfile[1] == "a" else False, verbosity=args.verbose, server_replay_nopop=args.server_replay_nopop, server_replay_ignore_content=args.server_replay_ignore_content, @@ -339,15 +340,15 @@ def basic_options(parser): action="store_const", dest="verbose", default=2, const=3, help="Increase log verbosity." ) - outfile = parser.add_mutually_exclusive_group() - outfile.add_argument( + streamfile = parser.add_mutually_exclusive_group() + streamfile.add_argument( "-w", "--wfile", - action="store", dest="outfile", type=lambda f: (f, "wb"), + action="store", dest="streamfile", type=lambda f: (f, "w"), help="Write flows to file." ) - outfile.add_argument( + streamfile.add_argument( "-a", "--afile", - action="store", dest="outfile", type=lambda f: (f, "ab"), + action="store", dest="streamfile", type=lambda f: (f, "a"), help="Append flows to file." ) parser.add_argument( diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py index 76545893c..6ad7f6562 100644 --- a/mitmproxy/tools/console/flowlist.py +++ b/mitmproxy/tools/console/flowlist.py @@ -393,13 +393,13 @@ class FlowListBox(urwid.ListBox): val = not self.master.options.order_reversed self.master.options.order_reversed = val elif key == "W": - if self.master.options.outfile: - self.master.options.outfile = None + if self.master.options.streamfile: + self.master.options.streamfile = None else: signals.status_prompt_path.send( self, prompt="Stream flows to", - callback= lambda path: self.master.options.update(outfile=(path, "ab")) + callback= lambda path: self.master.options.update(streamfile=path) ) else: return urwid.ListBox.keypress(self, size, key) diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py index e292cbd79..e34244937 100644 --- a/mitmproxy/tools/console/statusbar.py +++ b/mitmproxy/tools/console/statusbar.py @@ -223,8 +223,8 @@ class StatusBar(urwid.WidgetWrap): r.append(("heading_key", "s")) r.append("cripts:%s]" % len(self.master.options.scripts)) - if self.master.options.outfile: - r.append("[W:%s]" % self.master.options.outfile[0]) + if self.master.options.streamfile: + r.append("[W:%s]" % self.master.options.streamfile) return r diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index d283e5d5e..2f1fd4e39 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -111,15 +111,6 @@ class WebMaster(master.Master): "error" ) - if options.outfile: - err = self.start_stream_to_path( - options.outfile[0], - options.outfile[1] - ) - if err: - print("Stream file error: {}".format(err), file=sys.stderr) - sys.exit(1) - def _sig_add(self, view, flow): app.ClientConnection.broadcast( type="UPDATE_FLOWS", diff --git a/test/mitmproxy/addons/test_filestreamer.py b/test/mitmproxy/addons/test_filestreamer.py deleted file mode 100644 index 28094c436..000000000 --- a/test/mitmproxy/addons/test_filestreamer.py +++ /dev/null @@ -1,44 +0,0 @@ -from mitmproxy.test import tflow -from mitmproxy.test import tutils - -from .. import mastertest - -import os.path - -from mitmproxy.addons import filestreamer -from mitmproxy import master -from mitmproxy import io -from mitmproxy import options -from mitmproxy import proxy - - -class TestStream(mastertest.MasterTest): - def test_stream(self): - with tutils.tmpdir() as tdir: - p = os.path.join(tdir, "foo") - - def r(): - r = io.FlowReader(open(p, "rb")) - return list(r.stream()) - - o = options.Options( - outfile = (p, "wb") - ) - m = master.Master(o, proxy.DummyServer()) - sa = filestreamer.FileStreamer() - - m.addons.add(sa) - f = tflow.tflow(resp=True) - m.request(f) - m.response(f) - m.addons.remove(sa) - - assert r()[0].response - - m.options.outfile = (p, "ab") - - m.addons.add(sa) - f = tflow.tflow() - m.request(f) - m.addons.remove(sa) - assert not r()[1].response diff --git a/test/mitmproxy/addons/test_intercept.py b/test/mitmproxy/addons/test_intercept.py index efdac3744..a347f9abd 100644 --- a/test/mitmproxy/addons/test_intercept.py +++ b/test/mitmproxy/addons/test_intercept.py @@ -36,3 +36,8 @@ def test_simple(): f = tflow.tflow(resp=False) tctx.cycle(r, f) assert not f.intercepted + + f = tflow.tflow(resp=True) + f.reply._state = "handled" + r.response(f) + assert f.intercepted diff --git a/test/mitmproxy/addons/test_streamfile.py b/test/mitmproxy/addons/test_streamfile.py new file mode 100644 index 000000000..82a4345b4 --- /dev/null +++ b/test/mitmproxy/addons/test_streamfile.py @@ -0,0 +1,64 @@ +from mitmproxy.test import tflow +from mitmproxy.test import tutils +from mitmproxy.test import taddons + +import os.path +from mitmproxy import io +from mitmproxy import exceptions +from mitmproxy.tools import dump +from mitmproxy.addons import streamfile + + +def test_configure(): + sa = streamfile.StreamFile() + with taddons.context(options=dump.Options()) as tctx: + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") + tutils.raises( + exceptions.OptionsError, + tctx.configure, sa, streamfile=tdir + ) + tutils.raises( + "invalid filter", + tctx.configure, sa, streamfile=p, filtstr="~~" + ) + + +def rd(p): + x = io.FlowReader(open(p, "rb")) + return list(x.stream()) + + +def test_tcp(): + sa = streamfile.StreamFile() + with taddons.context() as tctx: + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") + tctx.configure(sa, streamfile=p) + + tt = tflow.ttcpflow() + sa.tcp_start(tt) + sa.tcp_end(tt) + tctx.configure(sa, streamfile=None) + assert rd(p) + + +def test_simple(): + sa = streamfile.StreamFile() + with taddons.context() as tctx: + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") + + tctx.configure(sa, streamfile=p) + + f = tflow.tflow(resp=True) + sa.request(f) + sa.response(f) + tctx.configure(sa, streamfile=None) + assert rd(p)[0].response + + tctx.configure(sa, streamfile=p, streamfile_append=True) + f = tflow.tflow() + sa.request(f) + tctx.configure(sa, streamfile=None) + assert not rd(p)[1].response diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py index aa3228e4e..e331637d9 100644 --- a/test/mitmproxy/test_dump.py +++ b/test/mitmproxy/test_dump.py @@ -2,7 +2,6 @@ from mitmproxy.test import tflow import os import io -import mitmproxy.io from mitmproxy.tools import dump from mitmproxy import exceptions from mitmproxy import proxy @@ -126,33 +125,6 @@ class TestDumpMaster(mastertest.MasterTest): f = self.cycle(m, b"content") assert f.request.headers["one"] == "two" - def test_write(self): - with tutils.tmpdir() as d: - p = os.path.join(d, "a") - self.dummy_cycle( - self.mkmaster(None, outfile=(p, "wb"), verbosity=0), 1, b"" - ) - assert len(list(mitmproxy.io.FlowReader(open(p, "rb")).stream())) == 1 - - def test_write_append(self): - with tutils.tmpdir() as d: - p = os.path.join(d, "a.append") - self.dummy_cycle( - self.mkmaster(None, outfile=(p, "wb"), verbosity=0), - 1, b"" - ) - self.dummy_cycle( - self.mkmaster(None, outfile=(p, "ab"), verbosity=0), - 1, b"" - ) - assert len(list(mitmproxy.io.FlowReader(open(p, "rb")).stream())) == 2 - - def test_write_err(self): - tutils.raises( - exceptions.OptionsError, - self.mkmaster, None, outfile = ("nonexistentdir/foo", "wb") - ) - def test_script(self): ret = self.dummy_cycle( self.mkmaster(