From 214498f01c86bd2c84d379880f5eeac521a55e66 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 14 Apr 2018 10:24:03 +1200 Subject: asyncio: adjust readfile.py addon for async --- mitmproxy/addons/readfile.py | 30 +++++++++++++++++------------- mitmproxy/tools/main.py | 1 - 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py index aaf02d01..f97e297b 100644 --- a/mitmproxy/addons/readfile.py +++ b/mitmproxy/addons/readfile.py @@ -1,3 +1,4 @@ +import asyncio import os.path import sys import typing @@ -17,12 +18,12 @@ class ReadFile: "Read flows from file." ) - def load_flows(self, fo: typing.IO[bytes]) -> int: + 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) + await ctx.master.load_flow(flow) cnt += 1 except (IOError, exceptions.FlowReadException) as e: if cnt: @@ -33,29 +34,32 @@ 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: + async def load_flows_from_path(self, path: str) -> int: if path == "-": - return self.load_flows(sys.stdin.buffer) + return await self.load_flows(sys.stdin.buffer) else: - return super().load_flows_from_path(path) + return await super().load_flows_from_path(path) diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 53c236bb..b424490a 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -150,7 +150,6 @@ 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, ) return {} -- cgit v1.2.3 From 8609de6f319b41b90398ed1e7184d59bbf31f4d3 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 14 Apr 2018 11:32:16 +1200 Subject: 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. --- mitmproxy/addons/readfile.py | 25 ++++++++- mitmproxy/tools/main.py | 1 + setup.py | 1 + test/mitmproxy/addons/test_readfile.py | 94 ++++++++++++++++++++-------------- 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 -- cgit v1.2.3 From 909a02ea0310d3b4ff822d4f3fbc07dd21ad36a1 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 14 Apr 2018 13:29:35 +1200 Subject: mitmdump: also set dumper_filter to default filter on startup Fixes #3051 --- mitmproxy/tools/main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 5f461419..87f1e9f3 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -152,6 +152,7 @@ def mitmdump(args=None): # pragma: no cover return dict( save_stream_filter=v, readfile_filter=v, + dumper_filter=v, ) return {} -- cgit v1.2.3