From b2c4f301cb834ecdf6e5b0063e86be877d3ece6d Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Thu, 14 Jul 2016 12:43:17 +1200
Subject: Stream to file -> addon

This commit also clarifies a confusion about the "outfile" attribute and its
use in testing in the mitmdump master.
---
 mitmproxy/builtins/__init__.py         |  2 +
 mitmproxy/builtins/stream.py           | 54 +++++++++++++++++++++
 mitmproxy/dump.py                      | 22 +++------
 mitmproxy/flow/master.py               | 28 -----------
 test/mitmproxy/builtins/test_stream.py | 46 ++++++++++++++++++
 test/mitmproxy/mastertest.py           | 13 +++--
 test/mitmproxy/test_dump.py            | 89 ++++++++++++++++++++--------------
 test/mitmproxy/test_flow.py            | 31 ------------
 8 files changed, 167 insertions(+), 118 deletions(-)
 create mode 100644 mitmproxy/builtins/stream.py
 create mode 100644 test/mitmproxy/builtins/test_stream.py

diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py
index 75326712..8021c20f 100644
--- a/mitmproxy/builtins/__init__.py
+++ b/mitmproxy/builtins/__init__.py
@@ -4,6 +4,7 @@ from mitmproxy.builtins import anticache
 from mitmproxy.builtins import anticomp
 from mitmproxy.builtins import stickyauth
 from mitmproxy.builtins import stickycookie
+from mitmproxy.builtins import stream
 
 
 def default_addons():
@@ -12,4 +13,5 @@ def default_addons():
         anticomp.AntiComp(),
         stickyauth.StickyAuth(),
         stickycookie.StickyCookie(),
+        stream.Stream(),
     ]
diff --git a/mitmproxy/builtins/stream.py b/mitmproxy/builtins/stream.py
new file mode 100644
index 00000000..821a71f1
--- /dev/null
+++ b/mitmproxy/builtins/stream.py
@@ -0,0 +1,54 @@
+from __future__ import absolute_import, print_function, division
+import os.path
+
+from mitmproxy import ctx
+from mitmproxy import exceptions
+from mitmproxy.flow import io
+
+
+class Stream:
+    def __init__(self):
+        self.stream = None
+
+    def start_stream_to_path(self, path, mode, filt):
+        path = os.path.expanduser(path)
+        try:
+            f = open(path, mode)
+        except IOError as v:
+            return str(v)
+        self.stream = io.FilteredFlowWriter(f, filt)
+
+    def configure(self, options):
+        # We're already streaming - stop the previous stream and restart
+        if self.stream:
+            self.done()
+
+        if options.outfile:
+            filt = None
+            if options.get("filtstr"):
+                filt = filt.parse(options.filtstr)
+                if not 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, filt)
+            if err:
+                raise exceptions.OptionsError(err)
+
+    def done(self):
+        if self.stream:
+            for flow in ctx.master.active_flows:
+                self.stream.add(flow)
+            self.stream.fo.close()
+            self.stream = None
+
+    def tcp_close(self, flow):
+        if self.stream:
+            self.stream.add(flow)
+
+    def response(self, flow):
+        if self.stream:
+            self.stream.add(flow)
diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py
index d9142a4d..296419db 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -52,16 +52,17 @@ class Options(options.Options):
         "replay_ignore_content",
         "replay_ignore_params",
         "replay_ignore_payload_params",
-        "replay_ignore_host"
+        "replay_ignore_host",
+
+        "tfile"
     ]
 
 
 class DumpMaster(flow.FlowMaster):
 
-    def __init__(self, server, options, outfile=None):
+    def __init__(self, server, options):
         flow.FlowMaster.__init__(self, options, server, flow.State())
         self.addons.add(*builtins.default_addons())
-        self.outfile = outfile
         self.o = options
         self.showhost = options.showhost
         self.replay_ignore_params = options.replay_ignore_params
@@ -82,15 +83,6 @@ class DumpMaster(flow.FlowMaster):
         else:
             self.filt = None
 
-        if options.outfile:
-            err = self.start_stream_to_path(
-                options.outfile[0],
-                options.outfile[1],
-                self.filt
-            )
-            if err:
-                raise DumpError(err)
-
         if options.replacements:
             for i in options.replacements:
                 self.replacehooks.add(*i)
@@ -163,7 +155,7 @@ class DumpMaster(flow.FlowMaster):
     def echo(self, text, indent=None, **style):
         if indent:
             text = self.indent(indent, text)
-        click.secho(text, file=self.outfile, **style)
+        click.secho(text, file=self.options.tfile, **style)
 
     def _echo_message(self, message):
         if self.options.flow_detail >= 2 and hasattr(message, "headers"):
@@ -312,8 +304,8 @@ class DumpMaster(flow.FlowMaster):
         if f.error:
             self.echo(" << {}".format(f.error.msg), bold=True, fg="red")
 
-        if self.outfile:
-            self.outfile.flush()
+        if self.options.tfile:
+            self.options.tfile.flush()
 
     def _process_flow(self, f):
         if self.filt and not f.match(self.filt):
diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py
index d67ee7cc..27ceee87 100644
--- a/mitmproxy/flow/master.py
+++ b/mitmproxy/flow/master.py
@@ -46,7 +46,6 @@ class FlowMaster(controller.Master):
         self.replay_ignore_content = None
         self.replay_ignore_host = False
 
-        self.stream = None
         self.apps = modules.AppRegistry()
 
     def start_app(self, host, port):
@@ -409,8 +408,6 @@ class FlowMaster(controller.Master):
         if not f.reply.acked:
             if self.client_playback:
                 self.client_playback.clear(f)
-        if self.stream:
-            self.stream.add(f)
         return f
 
     def handle_intercept(self, f):
@@ -471,33 +468,8 @@ class FlowMaster(controller.Master):
     @controller.handler
     def tcp_close(self, flow):
         self.active_flows.discard(flow)
-        if self.stream:
-            self.stream.add(flow)
         self.run_scripts("tcp_close", flow)
 
     def shutdown(self):
         super(FlowMaster, self).shutdown()
-
-        # Add all flows that are still active
-        if self.stream:
-            for flow in self.active_flows:
-                self.stream.add(flow)
-            self.stop_stream()
-
         self.unload_scripts()
-
-    def start_stream(self, fp, filt):
-        self.stream = io.FilteredFlowWriter(fp, filt)
-
-    def stop_stream(self):
-        self.stream.fo.close()
-        self.stream = None
-
-    def start_stream_to_path(self, path, mode="wb", filt=None):
-        path = os.path.expanduser(path)
-        try:
-            f = open(path, mode)
-            self.start_stream(f, filt)
-        except IOError as v:
-            return str(v)
-        self.stream_path = path
diff --git a/test/mitmproxy/builtins/test_stream.py b/test/mitmproxy/builtins/test_stream.py
new file mode 100644
index 00000000..54e4f7d9
--- /dev/null
+++ b/test/mitmproxy/builtins/test_stream.py
@@ -0,0 +1,46 @@
+from __future__ import absolute_import, print_function, division
+
+from .. import tutils, mastertest
+
+import os.path
+
+from mitmproxy.builtins import stream
+from mitmproxy.flow import master, FlowReader
+from mitmproxy.flow import state
+from mitmproxy import options
+
+
+class TestStream(mastertest.MasterTest):
+    def test_stream(self):
+        with tutils.tmpdir() as tdir:
+            p = os.path.join(tdir, "foo")
+
+            def r():
+                r = FlowReader(open(p, "rb"))
+                return list(r.stream())
+
+            s = state.State()
+            m = master.FlowMaster(
+                options.Options(
+                    outfile = (p, "wb")
+                ),
+                None,
+                s
+            )
+            sa = stream.Stream()
+
+            m.addons.add(sa)
+            f = tutils.tflow(resp=True)
+            self.invoke(m, "request", f)
+            self.invoke(m, "response", f)
+            m.addons.remove(sa)
+
+            assert r()[0].response
+
+            m.options.outfile = (p, "ab")
+
+            m.addons.add(sa)
+            f = tutils.tflow()
+            self.invoke(m, "request", f)
+            m.addons.remove(sa)
+            assert not r()[1].response
diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py
index 06854e25..9754d3a9 100644
--- a/test/mitmproxy/mastertest.py
+++ b/test/mitmproxy/mastertest.py
@@ -18,15 +18,14 @@ class MasterTest:
         l = proxy.Log("connect")
         l.reply = mock.MagicMock()
         master.log(l)
-        master.clientconnect(f.client_conn)
-        master.serverconnect(f.server_conn)
-        master.request(f)
+        self.invoke(master, "clientconnect", f.client_conn)
+        self.invoke(master, "clientconnect", f.client_conn)
+        self.invoke(master, "serverconnect", f.server_conn)
+        self.invoke(master, "request", f)
         if not f.error:
             f.response = models.HTTPResponse.wrap(netlib.tutils.tresp(content=content))
-            f.reply.acked = False
-            f = master.response(f)
-        f.client_conn.reply.acked = False
-        master.clientdisconnect(f.client_conn)
+            self.invoke(master, "response", f)
+        self.invoke(master, "clientdisconnect", f)
         return f
 
     def dummy_cycle(self, master, n, content):
diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py
index aa73b5a4..9686be84 100644
--- a/test/mitmproxy/test_dump.py
+++ b/test/mitmproxy/test_dump.py
@@ -4,31 +4,33 @@ from mitmproxy.exceptions import ContentViewException
 
 import netlib.tutils
 
-from mitmproxy import dump, flow, models
+from mitmproxy import dump, flow, models, exceptions
 from . import tutils, mastertest
 import mock
 
 
 def test_strfuncs():
-    o = dump.Options()
+    o = dump.Options(
+        tfile = StringIO(),
+        flow_detail = 0,
+    )
     m = dump.DumpMaster(None, o)
 
-    m.outfile = StringIO()
     m.o.flow_detail = 0
     m.echo_flow(tutils.tflow())
-    assert not m.outfile.getvalue()
+    assert not o.tfile.getvalue()
 
     m.o.flow_detail = 4
     m.echo_flow(tutils.tflow())
-    assert m.outfile.getvalue()
+    assert o.tfile.getvalue()
 
-    m.outfile = StringIO()
+    o.tfile = StringIO()
     m.echo_flow(tutils.tflow(resp=True))
-    assert "<<" in m.outfile.getvalue()
+    assert "<<" in o.tfile.getvalue()
 
-    m.outfile = StringIO()
+    o.tfile = StringIO()
     m.echo_flow(tutils.tflow(err=True))
-    assert "<<" in m.outfile.getvalue()
+    assert "<<" in o.tfile.getvalue()
 
     flow = tutils.tflow()
     flow.request = netlib.tutils.treq()
@@ -50,25 +52,32 @@ def test_strfuncs():
 def test_contentview(get_content_view):
     get_content_view.side_effect = ContentViewException(""), ("x", iter([]))
 
-    o = dump.Options(flow_detail=4, verbosity=3)
-    m = dump.DumpMaster(None, o, StringIO())
+    o = dump.Options(
+        flow_detail=4,
+        verbosity=3,
+        tfile=StringIO(),
+    )
+    m = dump.DumpMaster(None, o)
     m.echo_flow(tutils.tflow())
-    assert "Content viewer failed" in m.outfile.getvalue()
+    assert "Content viewer failed" in m.options.tfile.getvalue()
 
 
 class TestDumpMaster(mastertest.MasterTest):
     def dummy_cycle(self, master, n, content):
         mastertest.MasterTest.dummy_cycle(self, master, n, content)
-        return master.outfile.getvalue()
+        return master.options.tfile.getvalue()
 
     def mkmaster(self, filt, **options):
-        cs = StringIO()
         if "verbosity" not in options:
             options["verbosity"] = 0
         if "flow_detail" not in options:
             options["flow_detail"] = 0
-        o = dump.Options(filtstr=filt, **options)
-        return dump.DumpMaster(None, o, outfile=cs)
+        o = dump.Options(
+            filtstr=filt,
+            tfile=StringIO(),
+            **options
+        )
+        return dump.DumpMaster(None, o)
 
     def test_basic(self):
         for i in (1, 2, 3):
@@ -89,31 +98,33 @@ class TestDumpMaster(mastertest.MasterTest):
             )
 
     def test_error(self):
-        cs = StringIO()
-        o = dump.Options(flow_detail=1)
-        m = dump.DumpMaster(None, o, outfile=cs)
+        o = dump.Options(
+            tfile=StringIO(),
+            flow_detail=1
+        )
+        m = dump.DumpMaster(None, o)
         f = tutils.tflow(err=True)
         m.request(f)
         assert m.error(f)
-        assert "error" in cs.getvalue()
+        assert "error" in o.tfile.getvalue()
 
     def test_missing_content(self):
-        cs = StringIO()
-        o = dump.Options(flow_detail=3)
-        m = dump.DumpMaster(None, o, outfile=cs)
+        o = dump.Options(
+            flow_detail=3,
+            tfile=StringIO(),
+        )
+        m = dump.DumpMaster(None, o)
         f = tutils.tflow()
         f.request.content = None
         m.request(f)
         f.response = models.HTTPResponse.wrap(netlib.tutils.tresp())
         f.response.content = None
         m.response(f)
-        assert "content missing" in cs.getvalue()
+        assert "content missing" in o.tfile.getvalue()
 
     def test_replay(self):
-        cs = StringIO()
-
         o = dump.Options(server_replay=["nonexistent"], kill=True)
-        tutils.raises(dump.DumpError, dump.DumpMaster, None, o, outfile=cs)
+        tutils.raises(dump.DumpError, dump.DumpMaster, None, o)
 
         with tutils.tmpdir() as t:
             p = os.path.join(t, "rep")
@@ -122,7 +133,7 @@ class TestDumpMaster(mastertest.MasterTest):
             o = dump.Options(server_replay=[p], kill=True)
             o.verbosity = 0
             o.flow_detail = 0
-            m = dump.DumpMaster(None, o, outfile=cs)
+            m = dump.DumpMaster(None, o)
 
             self.cycle(m, b"content")
             self.cycle(m, b"content")
@@ -130,13 +141,13 @@ class TestDumpMaster(mastertest.MasterTest):
             o = dump.Options(server_replay=[p], kill=False)
             o.verbosity = 0
             o.flow_detail = 0
-            m = dump.DumpMaster(None, o, outfile=cs)
+            m = dump.DumpMaster(None, o)
             self.cycle(m, b"nonexistent")
 
             o = dump.Options(client_replay=[p], kill=False)
             o.verbosity = 0
             o.flow_detail = 0
-            m = dump.DumpMaster(None, o, outfile=cs)
+            m = dump.DumpMaster(None, o)
 
     def test_read(self):
         with tutils.tmpdir() as t:
@@ -172,20 +183,24 @@ class TestDumpMaster(mastertest.MasterTest):
         assert len(m.apps.apps) == 1
 
     def test_replacements(self):
-        cs = StringIO()
-        o = dump.Options(replacements=[(".*", "content", "foo")])
+        o = dump.Options(
+            replacements=[(".*", "content", "foo")],
+            tfile = StringIO(),
+        )
         o.verbosity = 0
         o.flow_detail = 0
-        m = dump.DumpMaster(None, o, outfile=cs)
+        m = dump.DumpMaster(None, o)
         f = self.cycle(m, b"content")
         assert f.request.content == b"foo"
 
     def test_setheader(self):
-        cs = StringIO()
-        o = dump.Options(setheaders=[(".*", "one", "two")])
+        o = dump.Options(
+            setheaders=[(".*", "one", "two")],
+            tfile=StringIO()
+        )
         o.verbosity = 0
         o.flow_detail = 0
-        m = dump.DumpMaster(None, o, outfile=cs)
+        m = dump.DumpMaster(None, o)
         f = self.cycle(m, b"content")
         assert f.request.headers["one"] == "two"
 
@@ -212,7 +227,7 @@ class TestDumpMaster(mastertest.MasterTest):
 
     def test_write_err(self):
         tutils.raises(
-            dump.DumpError,
+            exceptions.OptionsError,
             self.mkmaster, None, outfile = ("nonexistentdir/foo", "wb")
         )
 
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 91342e58..1a07f74d 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -1,5 +1,3 @@
-import os.path
-
 import mock
 import io
 
@@ -887,35 +885,6 @@ class TestFlowMaster:
         fm.process_new_request(f)
         assert "killed" in f.error.msg
 
-    def test_stream(self):
-        with tutils.tmpdir() as tdir:
-            p = os.path.join(tdir, "foo")
-
-            def read():
-                with open(p, "rb") as f:
-                    r = flow.FlowReader(f)
-                    return list(r.stream())
-
-            s = flow.State()
-            fm = flow.FlowMaster(None, None, s)
-            f = tutils.tflow(resp=True)
-
-            with open(p, "ab") as tmpfile:
-                fm.start_stream(tmpfile, None)
-                fm.request(f)
-                fm.response(f)
-                fm.stop_stream()
-
-            assert read()[0].response
-
-            with open(p, "ab") as tmpfile:
-                f = tutils.tflow()
-                fm.start_stream(tmpfile, None)
-                fm.request(f)
-                fm.shutdown()
-
-            assert not read()[1].response
-
 
 class TestRequest:
 
-- 
cgit v1.2.3