diff options
| -rw-r--r-- | .gitignore | 1 | ||||
| -rw-r--r-- | mitmproxy/addons/filestreamer.py | 21 | ||||
| -rw-r--r-- | mitmproxy/test/taddons.py | 18 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_filestreamer.py | 51 |
4 files changed, 69 insertions, 22 deletions
@@ -19,4 +19,3 @@ bower_components *.map sslkeylogfile.log .tox/ -.vscode diff --git a/mitmproxy/addons/filestreamer.py b/mitmproxy/addons/filestreamer.py index 92224d21..cb6f3762 100644 --- a/mitmproxy/addons/filestreamer.py +++ b/mitmproxy/addons/filestreamer.py @@ -8,6 +8,7 @@ from mitmproxy import io class FileStreamer: def __init__(self): self.stream = None + self.filt = None self.active_flows = set() # type: Set[flow.Flow] def start_stream_to_path(self, path, mode, flt): @@ -15,29 +16,27 @@ class FileStreamer: try: f = open(path, mode) except IOError as v: - return str(v) + raise exceptions.OptionsError(str(v)) self.stream = io.FilteredFlowWriter(f, flt) self.active_flows = set() def configure(self, options, updated): # We're already streaming - stop the previous stream and restart + if "filtstr" in updated: + if options.get("filtstr"): + self.filt = flowfilter.parse(options.filtstr) + if not self.filt: + raise exceptions.OptionsError( + "Invalid filter specification: %s" % options.filtstr + ) if "outfile" in updated: if self.stream: self.done() if options.outfile: - flt = None - if options.get("filtstr"): - flt = flowfilter.parse(options.filtstr) - if not flt: - raise exceptions.OptionsError( - "Invalid filter specification: %s" % options.filtstr - ) path, mode = options.outfile if mode not in ("wb", "ab"): raise exceptions.OptionsError("Invalid mode.") - err = self.start_stream_to_path(path, mode, flt) - if err: - raise exceptions.OptionsError(err) + self.start_stream_to_path(path, mode, self.filt) def tcp_start(self, flow): if self.stream: diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 3cba6762..7804b90d 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -1,7 +1,10 @@ +import contextlib + import mitmproxy.master import mitmproxy.options from mitmproxy import proxy from mitmproxy import events +from mitmproxy import exceptions class RecordingMaster(mitmproxy.master.Master): @@ -36,6 +39,15 @@ class context: self.wrapped = None return False + @contextlib.contextmanager + def _rollback(self, opts, updates): + old = opts._opts.copy() + try: + yield + except exceptions.OptionsError as e: + opts.__dict__["_opts"] = old + raise + def cycle(self, addon, f): """ Cycles the flow through the events for the flow. Stops if a reply @@ -55,6 +67,6 @@ class context: Options object with the given keyword arguments, then calls the configure method on the addon with the updated value. """ - for k, v in kwargs.items(): - setattr(self.options, k, v) - addon.configure(self.options, kwargs.keys()) + with self._rollback(self.options, kwargs): + self.options.update(**kwargs) + addon.configure(self.options, kwargs.keys()) diff --git a/test/mitmproxy/addons/test_filestreamer.py b/test/mitmproxy/addons/test_filestreamer.py index f86912fc..658a0aa8 100644 --- a/test/mitmproxy/addons/test_filestreamer.py +++ b/test/mitmproxy/addons/test_filestreamer.py @@ -3,19 +3,55 @@ from mitmproxy.test import tutils from mitmproxy.test import taddons import os.path -from mitmproxy.addons import filestreamer from mitmproxy import io +from mitmproxy import exceptions +from mitmproxy.tools import dump +from mitmproxy.addons import filestreamer + + +def test_configure(): + sa = filestreamer.FileStreamer() + with taddons.context(options=dump.Options()) as tctx: + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") + tutils.raises( + exceptions.OptionsError, + tctx.configure, sa, outfile=(tdir, "ab") + ) + tutils.raises( + "invalid filter", + tctx.configure, sa, outfile=(p, "ab"), filtstr="~~" + ) + tutils.raises( + "invalid mode", + tctx.configure, sa, outfile=(p, "xx") + ) + + +def rd(p): + x = io.FlowReader(open(p, "rb")) + return list(x.stream()) -def test_stream(): +def test_tcp(): sa = filestreamer.FileStreamer() with taddons.context() as tctx: with tutils.tmpdir() as tdir: p = os.path.join(tdir, "foo") + tctx.configure(sa, outfile=(p, "wb")) + + tt = tflow.ttcpflow() + sa.tcp_start(tt) + sa.tcp_end(tt) + tctx.configure(sa, outfile=None) + assert rd(p) - def r(): - x = io.FlowReader(open(p, "rb")) - return list(x.stream()) + +def test_simple(): + sa = filestreamer.FileStreamer() + with taddons.context() as tctx: + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") tctx.configure(sa, outfile=(p, "wb")) @@ -23,10 +59,11 @@ def test_stream(): sa.request(f) sa.response(f) tctx.configure(sa, outfile=None) - assert r()[0].response + assert rd(p)[0].response tctx.configure(sa, outfile=(p, "ab")) f = tflow.tflow() sa.request(f) tctx.configure(sa, outfile=None) - assert not r()[1].response + assert not rd(p)[1].response + |