readfile: add a readfile_filter option

* Add a readfile_filter option that filters flows on read.
* Adjust test suite for asyncio.
* Add asynctest as a dev dependency.

4 files changed, 81 insertions, 40 deletions

diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py
index f97e297b..2ae81fae 100644
--- a/mitmproxy/addons/readfile.py
+++ b/mitmproxy/addons/readfile.py
@@ -5,6 +5,7 @@ import typing
 
 from mitmproxy import ctx
 from mitmproxy import exceptions
+from mitmproxy import flowfilter
 from mitmproxy import io
 
 
@@ -12,17 +13,37 @@ 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 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():
+                if self.filter and not self.filter(flow):
+                    continue
                 await ctx.master.load_flow(flow)
                 cnt += 1
         except (IOError, exceptions.FlowReadException) as e:
@@ -59,7 +80,9 @@ class ReadFile:
 class ReadFileStdin(ReadFile):
     """Support the special case of "-" for reading from stdin"""
     async def load_flows_from_path(self, path: str) -> int:
-        if path == "-":
+        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 await super().load_flows_from_path(path)
diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py
index b424490a..5f461419 100644
--- a/mitmproxy/tools/main.py
+++ b/mitmproxy/tools/main.py
@@ -151,6 +151,7 @@ def mitmdump(args=None):  # pragma: no cover
             v = " ".join(args.filter_args)
             return dict(
                 save_stream_filter=v,
+                readfile_filter=v,
             )
         return {}
 
diff --git a/setup.py b/setup.py
index 3e914847..712ad5d2 100644
--- a/setup.py
+++ b/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",
diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py
index a0826a26..f7e0c5c5 100644
--- a/test/mitmproxy/addons/test_readfile.py
+++ b/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):
+    @asynctest.patch('sys.stdin')
+    @pytest.mark.asyncio
+    async def test_stdin(self, stdin, 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
+        with taddons.context(rf):
+            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
-            tctx.configure(rf, rfile="-")
-            with pytest.raises(exceptions.OptionsError):
-                rf.running()
+                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')
-    def test_normal(self, load_flow, tmpdir, data):
+    async def test_normal(self, load_flow, tmpdir, 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 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