Merge pull request #3056 from cortesi/readfile

readfile fixes

mitmproxy/addons/readfile.py

@@ -1,9 +1,11 @@
+import asyncio
 import os.path
 import sys
 import typing
 
 from mitmproxy import ctx
 from mitmproxy import exceptions
+from mitmproxy import flowfilter
 from mitmproxy import io
 
 
@@ -11,18 +13,38 @@ class ReadFile:
     """
         An addon that handles reading from file on startup.
     """
+    def __init__(self):
+        self.filter = None
+
     def load(self, loader):
         loader.add_option(
             "rfile", typing.Optional[str], None,
             "Read flows from file."
         )
+        loader.add_option(
+            "readfile_filter", typing.Optional[str], None,
+            "Read only matching flows."
+        )
 
-    def load_flows(self, fo: typing.IO[bytes]) -> int:
+    def configure(self, updated):
+        if "readfile_filter" in updated:
+            filt = None
+            if ctx.options.readfile_filter:
+                filt = flowfilter.parse(ctx.options.readfile_filter)
+                if not filt:
+                    raise exceptions.OptionsError(
+                        "Invalid readfile filter: %s" % ctx.options.readfile_filter
+                    )
+            self.filter = filt
+
+    async def load_flows(self, fo: typing.IO[bytes]) -> int:
         cnt = 0
         freader = io.FlowReader(fo)
         try:
             for flow in freader.stream():
-                ctx.master.load_flow(flow)
+                if self.filter and not self.filter(flow):
+                    continue
+                await ctx.master.load_flow(flow)
                 cnt += 1
         except (IOError, exceptions.FlowReadException) as e:
             if cnt:
@@ -33,29 +55,34 @@ class ReadFile:
         else:
             return cnt
 
-    def load_flows_from_path(self, path: str) -> int:
+    async def load_flows_from_path(self, path: str) -> int:
         path = os.path.expanduser(path)
         try:
             with open(path, "rb") as f:
-                return self.load_flows(f)
+                return await self.load_flows(f)
         except IOError as e:
             ctx.log.error("Cannot load flows: {}".format(e))
             raise exceptions.FlowReadException(str(e)) from e
 
+    async def doread(self, rfile):
+        try:
+            await self.load_flows_from_path(ctx.options.rfile)
+        except exceptions.FlowReadException as e:
+            raise exceptions.OptionsError(e) from e
+        finally:
+            ctx.master.addons.trigger("processing_complete")
+
     def running(self):
         if ctx.options.rfile:
-            try:
-                self.load_flows_from_path(ctx.options.rfile)
-            except exceptions.FlowReadException as e:
-                raise exceptions.OptionsError(e) from e
-            finally:
-                ctx.master.addons.trigger("processing_complete")
+            asyncio.get_event_loop().create_task(self.doread(ctx.options.rfile))
 
 
 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)
+    async def load_flows_from_path(self, path: str) -> int:
+        if path == "-":  # pragma: no cover
+            # Need to think about how to test this. This function is scheduled
+            # onto the event loop, where a sys.stdin mock has no effect.
+            return await self.load_flows(sys.stdin.buffer)
         else:
-            return super().load_flows_from_path(path)
+            return await super().load_flows_from_path(path)

mitmproxy/tools/main.py

@@ -150,8 +150,9 @@ def mitmdump(args=None):  # pragma: no cover
         if args.filter_args:
             v = " ".join(args.filter_args)
             return dict(
-                view_filter=v,
                 save_stream_filter=v,
+                readfile_filter=v,
+                dumper_filter=v,
             )
         return {}
 

setup.py

@@ -85,6 +85,7 @@ setup(
             "pydivert>=2.0.3,<2.2",
         ],
         'dev': [
+            "asynctest>=0.12.0",
             "flake8>=3.5, <3.6",
             "Flask>=0.10.1, <0.13",
             "mypy>=0.580,<0.581",

test/mitmproxy/addons/test_readfile.py

@@ -1,7 +1,9 @@
+import asyncio
 import io
 from unittest import mock
 
 import pytest
+import asynctest
 
 import mitmproxy.io
 from mitmproxy import exceptions
@@ -38,69 +40,83 @@ def corrupt_data():
 
 
 class TestReadFile:
-    @mock.patch('mitmproxy.master.Master.load_flow')
-    def test_configure(self, mck, tmpdir, data, corrupt_data):
+    def test_configure(self):
+        rf = readfile.ReadFile()
+        with taddons.context() as tctx:
+            tctx.configure(rf, readfile_filter="~q")
+            with pytest.raises(Exception, match="Invalid readfile filter"):
+                tctx.configure(rf, readfile_filter="~~")
+
+    @pytest.mark.asyncio
+    async def test_read(self, tmpdir, data, corrupt_data):
         rf = readfile.ReadFile()
         with taddons.context(rf) 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
+            with asynctest.patch('mitmproxy.master.Master.load_flow') as mck:
+                tf.write(data.getvalue())
+                tctx.configure(
+                    rf,
+                    rfile = str(tf),
+                    readfile_filter = ".*"
+                )
+                assert not mck.awaited
+                rf.running()
+                await asyncio.sleep(0)
+                assert mck.awaited
 
             tf.write(corrupt_data.getvalue())
             tctx.configure(rf, rfile=str(tf))
-            with pytest.raises(exceptions.OptionsError):
-                rf.running()
+            rf.running()
+            assert await tctx.master.await_log("corrupted")
 
     @pytest.mark.asyncio
     async def test_corrupt(self, corrupt_data):
         rf = readfile.ReadFile()
         with taddons.context(rf) as tctx:
-            with mock.patch('mitmproxy.master.Master.load_flow') as mck:
-                with pytest.raises(exceptions.FlowReadException):
-                    rf.load_flows(io.BytesIO(b"qibble"))
-                assert not mck.called
+            with pytest.raises(exceptions.FlowReadException):
+                await rf.load_flows(io.BytesIO(b"qibble"))
 
-                tctx.master.clear()
-                with pytest.raises(exceptions.FlowReadException):
-                    rf.load_flows(corrupt_data)
-                assert await tctx.master.await_log("file corrupted")
-                assert mck.called
+            tctx.master.clear()
+            with pytest.raises(exceptions.FlowReadException):
+                await rf.load_flows(corrupt_data)
+            assert await tctx.master.await_log("file corrupted")
 
     @pytest.mark.asyncio
-    async def test_nonexisting_file(self):
+    async def test_nonexistent_file(self):
         rf = readfile.ReadFile()
         with taddons.context(rf) as tctx:
             with pytest.raises(exceptions.FlowReadException):
-                rf.load_flows_from_path("nonexistent")
+                await rf.load_flows_from_path("nonexistent")
             assert await tctx.master.await_log("nonexistent")
 
 
 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(rf) 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_normal(self, load_flow, tmpdir, data):
+    @asynctest.patch('sys.stdin')
+    @pytest.mark.asyncio
+    async def test_stdin(self, stdin, data, corrupt_data):
         rf = readfile.ReadFileStdin()
         with taddons.context(rf):
-            tfile = tmpdir.join("tfile")
-            tfile.write(data.getvalue())
-            rf.load_flows_from_path(str(tfile))
-            assert load_flow.called
+            with asynctest.patch('mitmproxy.master.Master.load_flow') as mck:
+                stdin.buffer = data
+                assert not mck.awaited
+                await rf.load_flows(stdin.buffer)
+                assert mck.awaited
+
+                stdin.buffer = corrupt_data
+                with pytest.raises(exceptions.FlowReadException):
+                    await rf.load_flows(stdin.buffer)
+
+    @pytest.mark.asyncio
+    @mock.patch('mitmproxy.master.Master.load_flow')
+    async def test_normal(self, load_flow, tmpdir, data):
+        rf = readfile.ReadFileStdin()
+        with taddons.context(rf) as tctx:
+            tf = tmpdir.join("tfile")
+            with asynctest.patch('mitmproxy.master.Master.load_flow') as mck:
+                tf.write(data.getvalue())
+                tctx.configure(rf, rfile=str(tf))
+                assert not mck.awaited
+                rf.running()
+                await asyncio.sleep(0)
+                assert mck.awaited