Merge pull request #1706 from cortesi/taddons

addons tests and bugfixes

13 files changed, 119 insertions, 132 deletions

diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py
index c5d40525..d2b50c35 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 50b52214..d418952c 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
index 031b44ab..377f277d 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 1db9f0f0..49791400 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 3cba6762..7804b90d 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 e4b29d0f..debe6db9 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 76545893..6ad7f656 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 e292cbd7..e3424493 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 d283e5d5..2f1fd4e3 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 28094c43..00000000
--- 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 efdac374..a347f9ab 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 00000000..82a4345b
--- /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 aa3228e4..e331637d 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(