From 2b500f234fc8cf5d7922218ac1fb12e9874c2f61 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Mon, 24 Apr 2017 13:39:59 +0200 Subject: typecheck: add support for typing.Any --- mitmproxy/utils/typecheck.py | 2 ++ test/mitmproxy/utils/test_typecheck.py | 6 ++++++ 2 files changed, 8 insertions(+) diff --git a/mitmproxy/utils/typecheck.py b/mitmproxy/utils/typecheck.py index 5df4ea4b..628ea642 100644 --- a/mitmproxy/utils/typecheck.py +++ b/mitmproxy/utils/typecheck.py @@ -68,5 +68,7 @@ def check_type(name: str, value: typing.Any, typeinfo: typing.Any) -> None: return else: raise e + elif typename.startswith("typing.Any"): + return elif not isinstance(value, typeinfo): raise e diff --git a/test/mitmproxy/utils/test_typecheck.py b/test/mitmproxy/utils/test_typecheck.py index d99a914f..fd0c6e0c 100644 --- a/test/mitmproxy/utils/test_typecheck.py +++ b/test/mitmproxy/utils/test_typecheck.py @@ -79,3 +79,9 @@ def test_check_io(): typecheck.check_type("foo", io.StringIO(), typing.IO[str]) with pytest.raises(TypeError): typecheck.check_type("foo", "foo", typing.IO[str]) + + +def test_check_any(): + typecheck.check_type("foo", 42, typing.Any) + typecheck.check_type("foo", object(), typing.Any) + typecheck.check_type("foo", None, typing.Any) -- cgit v1.2.3 From b3a11433389b47c913cfcd3c56a72d39ff6b93ef Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Mon, 24 Apr 2017 14:10:47 +0200 Subject: integrate readstdin into readfile addon --- mitmproxy/addons/readfile.py | 44 ++++++++++------ mitmproxy/addons/readstdin.py | 26 --------- mitmproxy/tools/dump.py | 3 +- test/mitmproxy/addons/test_readfile.py | 93 ++++++++++++++++++++++----------- test/mitmproxy/addons/test_readstdin.py | 53 ------------------- 5 files changed, 92 insertions(+), 127 deletions(-) delete mode 100644 mitmproxy/addons/readstdin.py delete mode 100644 test/mitmproxy/addons/test_readstdin.py 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 -- cgit v1.2.3 From ca2827886a97de88f4ab2937e71588fc9320ba4d Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Wed, 26 Apr 2017 11:12:14 +0200 Subject: separate reading from stdin into its own addon --- mitmproxy/addons/__init__.py | 2 - mitmproxy/addons/readfile.py | 32 +++++---- mitmproxy/tools/console/master.py | 10 ++- mitmproxy/tools/dump.py | 3 +- mitmproxy/tools/web/master.py | 2 + test/mitmproxy/addons/test_readfile.py | 117 ++++++++++++++++++--------------- 6 files changed, 94 insertions(+), 72 deletions(-) diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index b4367d78..7a45106c 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -8,7 +8,6 @@ from mitmproxy.addons import disable_h2c from mitmproxy.addons import onboarding from mitmproxy.addons import proxyauth from mitmproxy.addons import replace -from mitmproxy.addons import readfile from mitmproxy.addons import script from mitmproxy.addons import serverplayback from mitmproxy.addons import setheaders @@ -38,6 +37,5 @@ def default_addons(): stickycookie.StickyCookie(), streambodies.StreamBodies(), streamfile.StreamFile(), - readfile.ReadFile(), upstream_auth.UpstreamAuth(), ] diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py index e83c9993..05b6c309 100644 --- a/mitmproxy/addons/readfile.py +++ b/mitmproxy/addons/readfile.py @@ -1,17 +1,17 @@ import os.path -import typing - import sys +import typing from mitmproxy import ctx -from mitmproxy import io from mitmproxy import exceptions +from mitmproxy import io class ReadFile: """ An addon that handles reading from file on startup. """ + def load_flows(self, fo: typing.IO[bytes]) -> int: cnt = 0 freader = io.FlowReader(fo) @@ -29,16 +29,13 @@ class ReadFile: 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 + 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: @@ -48,3 +45,12 @@ class ReadFile: raise exceptions.OptionsError(e) from e finally: ctx.master.addons.trigger("processing_complete") + + +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) + else: + return super().load_flows_from_path(path) diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index c1d584ac..e7a2c6ae 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -17,8 +17,9 @@ from mitmproxy import exceptions from mitmproxy import master from mitmproxy import io from mitmproxy import log -from mitmproxy.addons import view from mitmproxy.addons import intercept +from mitmproxy.addons import readfile +from mitmproxy.addons import view from mitmproxy.tools.console import flowlist from mitmproxy.tools.console import flowview from mitmproxy.tools.console import grideditor @@ -91,7 +92,12 @@ class ConsoleMaster(master.Master): signals.sig_add_log.connect(self.sig_add_log) self.addons.add(Logger()) self.addons.add(*addons.default_addons()) - self.addons.add(intercept.Intercept(), self.view, UnsupportedLog()) + self.addons.add( + intercept.Intercept(), + self.view, + UnsupportedLog(), + readfile.ReadFile(), + ) def sigint_handler(*args, **kwargs): self.prompt_for_exit() diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py index a4c9998b..4d0ccf4b 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, keepserving +from mitmproxy.addons import dumper, termlog, termstatus, keepserving, readfile class ErrorCheck: @@ -31,5 +31,6 @@ class DumpMaster(master.Master): self.addons.add(dumper.Dumper()) self.addons.add( keepserving.KeepServing(), + readfile.ReadFileStdin(), self.errorcheck ) diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index 0db5a09f..c09fe0a2 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -7,6 +7,7 @@ from mitmproxy import log from mitmproxy import master from mitmproxy.addons import eventstore from mitmproxy.addons import intercept +from mitmproxy.addons import readfile from mitmproxy.addons import termlog from mitmproxy.addons import view from mitmproxy.addons import termstatus @@ -32,6 +33,7 @@ class WebMaster(master.Master): self.addons.add(*addons.default_addons()) self.addons.add( intercept.Intercept(), + readfile.ReadFile(), self.view, self.events, ) diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py index 689d9779..813aa10e 100644 --- a/test/mitmproxy/addons/test_readfile.py +++ b/test/mitmproxy/addons/test_readfile.py @@ -37,59 +37,68 @@ def corrupt_data(): return f -@mock.patch('mitmproxy.master.Master.load_flow') -def test_configure(mck, tmpdir, data, corrupt_data): - rf = readfile.ReadFile() - with taddons.context() 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 - - tf.write(corrupt_data.getvalue()) - tctx.configure(rf, rfile=str(tf)) - with pytest.raises(exceptions.OptionsError): +class TestReadFile: + @mock.patch('mitmproxy.master.Master.load_flow') + def test_configure(self, mck, tmpdir, data, corrupt_data): + rf = readfile.ReadFile() + with taddons.context() as tctx: + tf = tmpdir.join("tfile") + + tf.write(data.getvalue()) + tctx.configure(rf, rfile=str(tf)) + assert not mck.called rf.running() - - -@mock.patch('mitmproxy.master.Master.load_flow') -@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): + assert mck.called + + 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_corrupt(self, mck, corrupt_data): + rf = readfile.ReadFile() + with taddons.context() as tctx: + with pytest.raises(exceptions.FlowReadException): + rf.load_flows(io.BytesIO(b"qibble")) + assert not mck.called + assert len(tctx.master.logs) == 1 + + with pytest.raises(exceptions.FlowReadException): + rf.load_flows(corrupt_data) + assert mck.called + assert len(tctx.master.logs) == 2 + + def test_nonexisting_file(self): + 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 + + +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() as tctx: + stdin.buffer = data + tctx.configure(rf, rfile="-") + assert not load_flow.called 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(io.BytesIO(b"qibble")) - assert not mck.called - assert len(tctx.master.logs) == 1 - - with pytest.raises(exceptions.FlowReadException): - 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 + 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): + rf = readfile.ReadFileStdin() + with taddons.context(): + tfile = tmpdir.join("tfile") + tfile.write(data.getvalue()) + rf.load_flows_from_path(str(tfile)) + assert load_flow.called -- cgit v1.2.3