diff options
| -rw-r--r-- | mitmproxy/addons/readfile.py | 44 | ||||
| -rw-r--r-- | mitmproxy/addons/readstdin.py | 26 | ||||
| -rw-r--r-- | mitmproxy/tools/dump.py | 3 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_readfile.py | 93 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_readstdin.py | 53 |
5 files changed, 92 insertions, 127 deletions
diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py index 949da15d..e83c9993 100644 --- a/mitmproxy/addons/readfile.py +++ b/mitmproxy/addons/readfile.py @@ -1,4 +1,7 @@ import os.path +import typing + +import sys from mitmproxy import ctx from mitmproxy import io @@ -9,30 +12,39 @@ class ReadFile: """ An addon that handles reading from file on startup. """ - def load_flows_file(self, path: str) -> int: - path = os.path.expanduser(path) + def load_flows(self, fo: typing.IO[bytes]) -> int: cnt = 0 + freader = io.FlowReader(fo) try: - with open(path, "rb") as f: - freader = io.FlowReader(f) - for i in freader.stream(): - cnt += 1 - ctx.master.load_flow(i) - return cnt - except (IOError, exceptions.FlowReadException) as v: + for flow in freader.stream(): + ctx.master.load_flow(flow) + cnt += 1 + except (IOError, exceptions.FlowReadException) as e: if cnt: - ctx.log.warn( - "Flow file corrupted - loaded %i flows." % cnt, - ) + ctx.log.warn("Flow file corrupted - loaded %i flows." % cnt) else: ctx.log.error("Flow file corrupted.") - raise exceptions.FlowReadException(v) + raise exceptions.FlowReadException(str(e)) from e + else: + return cnt + + def load_flows_from_path(self, path: str) -> int: + if path == "-": + return self.load_flows(sys.stdin.buffer) + else: + path = os.path.expanduser(path) + try: + with open(path, "rb") as f: + return self.load_flows(f) + except IOError as e: + ctx.log.error("Cannot load flows: {}".format(e)) + raise exceptions.FlowReadException(str(e)) from e def running(self): if ctx.options.rfile: try: - self.load_flows_file(ctx.options.rfile) - except exceptions.FlowReadException as v: - raise exceptions.OptionsError(v) + 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") diff --git a/mitmproxy/addons/readstdin.py b/mitmproxy/addons/readstdin.py deleted file mode 100644 index 93a99f01..00000000 --- a/mitmproxy/addons/readstdin.py +++ /dev/null @@ -1,26 +0,0 @@ -from mitmproxy import ctx -from mitmproxy import io -from mitmproxy import exceptions -import sys - - -class ReadStdin: - """ - An addon that reads from stdin if we're not attached to (someting like) - a tty. - """ - def running(self, stdin = sys.stdin): - if not stdin.isatty(): - ctx.log.info("Reading from stdin") - try: - stdin.buffer.read(0) - except Exception as e: - ctx.log.warn("Cannot read from stdin: {}".format(e)) - return - freader = io.FlowReader(stdin.buffer) - try: - for i in freader.stream(): - ctx.master.load_flow(i) - except exceptions.FlowReadException as e: - ctx.log.error("Error reading from stdin: %s" % e) - ctx.master.addons.trigger("processing_complete") diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py index 6329f6b7..a4c9998b 100644 --- a/mitmproxy/tools/dump.py +++ b/mitmproxy/tools/dump.py @@ -1,7 +1,7 @@ from mitmproxy import addons from mitmproxy import options from mitmproxy import master -from mitmproxy.addons import dumper, termlog, termstatus, readstdin, keepserving +from mitmproxy.addons import dumper, termlog, termstatus, keepserving class ErrorCheck: @@ -30,7 +30,6 @@ class DumpMaster(master.Master): if with_dumper: self.addons.add(dumper.Dumper()) self.addons.add( - readstdin.ReadStdin(), keepserving.KeepServing(), self.errorcheck ) diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py index b30c147b..689d9779 100644 --- a/test/mitmproxy/addons/test_readfile.py +++ b/test/mitmproxy/addons/test_readfile.py @@ -1,62 +1,95 @@ +import io +from unittest import mock + +import pytest + +import mitmproxy.io +from mitmproxy import exceptions from mitmproxy.addons import readfile from mitmproxy.test import taddons from mitmproxy.test import tflow -from mitmproxy import io -from mitmproxy import exceptions -from unittest import mock -import pytest +@pytest.fixture +def data(): + f = io.BytesIO() -def write_data(path, corrupt=False): - with open(path, "wb") as tf: - w = io.FlowWriter(tf) - for i in range(3): - f = tflow.tflow(resp=True) - w.add(f) - for i in range(3): - f = tflow.tflow(err=True) - w.add(f) - f = tflow.ttcpflow() - w.add(f) - f = tflow.ttcpflow(err=True) - w.add(f) - if corrupt: - tf.write(b"flibble") + w = mitmproxy.io.FlowWriter(f) + flows = [ + tflow.tflow(resp=True), + tflow.tflow(err=True), + tflow.ttcpflow(), + tflow.ttcpflow(err=True) + ] + for flow in flows: + w.add(flow) + f.seek(0) + return f + + +@pytest.fixture +def corrupt_data(): + f = data() + f.seek(0, io.SEEK_END) + f.write(b"qibble") + f.seek(0) + return f -@mock.patch('mitmproxy.master.Master.load_flow') -def test_configure(mck, tmpdir): +@mock.patch('mitmproxy.master.Master.load_flow') +def test_configure(mck, tmpdir, data, corrupt_data): rf = readfile.ReadFile() with taddons.context() as tctx: - tf = str(tmpdir.join("tfile")) - write_data(tf) + tf = tmpdir.join("tfile") + + tf.write(data.getvalue()) tctx.configure(rf, rfile=str(tf)) assert not mck.called rf.running() assert mck.called - write_data(tf, corrupt=True) + tf.write(corrupt_data.getvalue()) tctx.configure(rf, rfile=str(tf)) with pytest.raises(exceptions.OptionsError): rf.running() @mock.patch('mitmproxy.master.Master.load_flow') -def test_corruption(mck, tmpdir): +@mock.patch('sys.stdin') +def test_configure_stdin(stdin, load_flow, data, corrupt_data): + rf = readfile.ReadFile() + with taddons.context() 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_corrupt(mck, corrupt_data): rf = readfile.ReadFile() with taddons.context() as tctx: with pytest.raises(exceptions.FlowReadException): - rf.load_flows_file("nonexistent") + rf.load_flows(io.BytesIO(b"qibble")) assert not mck.called assert len(tctx.master.logs) == 1 - tfc = str(tmpdir.join("tfile")) - write_data(tfc, corrupt=True) - with pytest.raises(exceptions.FlowReadException): - rf.load_flows_file(tfc) + rf.load_flows(corrupt_data) assert mck.called assert len(tctx.master.logs) == 2 + + +def test_nonexisting_file(): + rf = readfile.ReadFile() + with taddons.context() as tctx: + with pytest.raises(exceptions.FlowReadException): + rf.load_flows_from_path("nonexistent") + assert len(tctx.master.logs) == 1 diff --git a/test/mitmproxy/addons/test_readstdin.py b/test/mitmproxy/addons/test_readstdin.py deleted file mode 100644 index 76b01f4f..00000000 --- a/test/mitmproxy/addons/test_readstdin.py +++ /dev/null @@ -1,53 +0,0 @@ - -import io -from mitmproxy.addons import readstdin -from mitmproxy.test import taddons -from mitmproxy.test import tflow -import mitmproxy.io -from unittest import mock - - -def gen_data(corrupt=False): - tf = io.BytesIO() - w = mitmproxy.io.FlowWriter(tf) - for i in range(3): - f = tflow.tflow(resp=True) - w.add(f) - for i in range(3): - f = tflow.tflow(err=True) - w.add(f) - f = tflow.ttcpflow() - w.add(f) - f = tflow.ttcpflow(err=True) - w.add(f) - if corrupt: - tf.write(b"flibble") - tf.seek(0) - return tf - - -class mStdin: - def __init__(self, d): - self.buffer = d - - def isatty(self): - return False - - -@mock.patch('mitmproxy.master.Master.load_flow') -def test_read(m, tmpdir): - rf = readstdin.ReadStdin() - with taddons.context() as tctx: - assert not m.called - rf.running(stdin=mStdin(gen_data())) - assert m.called - - rf.running(stdin=mStdin(None)) - assert tctx.master.logs - tctx.master.clear() - - m.reset_mock() - assert not m.called - rf.running(stdin=mStdin(gen_data(corrupt=True))) - assert m.called - assert tctx.master.logs |