From b7afcb5dc2007ce451161e513d3e55d4cc91379e Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 27 Apr 2017 15:58:54 +1200 Subject: addons.streamfile -> addons.save Options: streamfile -> save_stream_file streamfile_filter -> save_stream_filter --- mitmproxy/addons/__init__.py | 4 +- mitmproxy/addons/save.py | 71 ++++++++++++++++++++++++++++++++ mitmproxy/addons/streamfile.py | 71 -------------------------------- mitmproxy/options.py | 6 +-- mitmproxy/test/taddons.py | 2 +- mitmproxy/tools/cmdline.py | 4 +- mitmproxy/tools/console/flowlist.py | 6 +-- mitmproxy/tools/console/statusbar.py | 4 +- mitmproxy/tools/main.py | 2 +- test/mitmproxy/addons/test_save.py | 62 ++++++++++++++++++++++++++++ test/mitmproxy/addons/test_streamfile.py | 62 ---------------------------- 11 files changed, 147 insertions(+), 147 deletions(-) create mode 100644 mitmproxy/addons/save.py delete mode 100644 mitmproxy/addons/streamfile.py create mode 100644 test/mitmproxy/addons/test_save.py delete mode 100644 test/mitmproxy/addons/test_streamfile.py diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index 7a45106c..e87f2cbd 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -14,7 +14,7 @@ from mitmproxy.addons import setheaders from mitmproxy.addons import stickyauth from mitmproxy.addons import stickycookie from mitmproxy.addons import streambodies -from mitmproxy.addons import streamfile +from mitmproxy.addons import save from mitmproxy.addons import upstream_auth @@ -36,6 +36,6 @@ def default_addons(): stickyauth.StickyAuth(), stickycookie.StickyCookie(), streambodies.StreamBodies(), - streamfile.StreamFile(), + save.Save(), upstream_auth.UpstreamAuth(), ] diff --git a/mitmproxy/addons/save.py b/mitmproxy/addons/save.py new file mode 100644 index 00000000..9cf65750 --- /dev/null +++ b/mitmproxy/addons/save.py @@ -0,0 +1,71 @@ +import os.path + +from mitmproxy import exceptions +from mitmproxy import flowfilter +from mitmproxy import io +from mitmproxy import ctx + + +class Save: + 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): + path = os.path.expanduser(path) + try: + f = open(path, mode) + except IOError as v: + raise exceptions.OptionsError(str(v)) + self.stream = io.FilteredFlowWriter(f, flt) + self.active_flows = set() + + def configure(self, updated): + # We're already streaming - stop the previous stream and restart + if "save_stream_filter" in updated: + if ctx.options.save_stream_filter: + self.filt = flowfilter.parse(ctx.options.save_stream_filter) + if not self.filt: + raise exceptions.OptionsError( + "Invalid filter specification: %s" % ctx.options.save_stream_filter + ) + else: + self.filt = None + if "save_stream_file" in updated: + if self.stream: + self.done() + if ctx.options.save_stream_file: + if ctx.options.save_stream_file.startswith("+"): + path = ctx.options.save_stream_file[1:] + mode = "ab" + else: + path = ctx.options.save_stream_file + mode = "wb" + self.start_stream_to_path(path, mode, self.filt) + + def tcp_start(self, flow): + if self.stream: + self.active_flows.add(flow) + + def tcp_end(self, flow): + if self.stream: + self.stream.add(flow) + self.active_flows.discard(flow) + + def response(self, flow): + if self.stream: + self.stream.add(flow) + self.active_flows.discard(flow) + + def request(self, flow): + if self.stream: + self.active_flows.add(flow) + + def done(self): + if self.stream: + for flow in self.active_flows: + self.stream.add(flow) + self.active_flows = set([]) + self.stream.fo.close() + self.stream = None diff --git a/mitmproxy/addons/streamfile.py b/mitmproxy/addons/streamfile.py deleted file mode 100644 index fde5a1c5..00000000 --- a/mitmproxy/addons/streamfile.py +++ /dev/null @@ -1,71 +0,0 @@ -import os.path - -from mitmproxy import exceptions -from mitmproxy import flowfilter -from mitmproxy import io -from mitmproxy import ctx - - -class StreamFile: - 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): - path = os.path.expanduser(path) - try: - f = open(path, mode) - except IOError as v: - raise exceptions.OptionsError(str(v)) - self.stream = io.FilteredFlowWriter(f, flt) - self.active_flows = set() - - def configure(self, updated): - # We're already streaming - stop the previous stream and restart - if "streamfile_filter" in updated: - if ctx.options.streamfile_filter: - self.filt = flowfilter.parse(ctx.options.streamfile_filter) - if not self.filt: - raise exceptions.OptionsError( - "Invalid filter specification: %s" % ctx.options.streamfile_filter - ) - else: - self.filt = None - if "streamfile" in updated: - if self.stream: - self.done() - if ctx.options.streamfile: - if ctx.options.streamfile.startswith("+"): - path = ctx.options.streamfile[1:] - mode = "ab" - else: - path = ctx.options.streamfile - mode = "wb" - self.start_stream_to_path(path, mode, self.filt) - - def tcp_start(self, flow): - if self.stream: - self.active_flows.add(flow) - - def tcp_end(self, flow): - if self.stream: - self.stream.add(flow) - self.active_flows.discard(flow) - - def response(self, flow): - if self.stream: - self.stream.add(flow) - self.active_flows.discard(flow) - - def request(self, flow): - if self.stream: - self.active_flows.add(flow) - - def done(self): - if self.stream: - for flow in self.active_flows: - self.stream.add(flow) - self.active_flows = set([]) - self.stream.fo.close() - self.stream = None diff --git a/mitmproxy/options.py b/mitmproxy/options.py index 8f8c1484..e477bed5 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -159,11 +159,11 @@ class Options(optmanager.OptManager): choices = [i.name.lower() for i in contentviews.views] ) self.add_option( - "streamfile", Optional[str], None, - "Write flows to file. Prefix path with + to append." + "save_stream_file", Optional[str], None, + "Stream flows to file as they arrive. Prefix path with + to append." ) self.add_option( - "streamfile_filter", Optional[str], None, + "save_stream_filter", Optional[str], None, "Filter which flows are written to file." ) self.add_option( diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 9e2c9838..5680e847 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -130,7 +130,7 @@ class context: def command(self, func, *args): """ - Invoke a command function within a command context, mimicing the actual command environment. + Invoke a command function with a list of string arguments within a command context, mimicing the actual command environment. """ cmd = command.Command(self.master.commands, "test.command", func) return cmd.call(args) diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index 2714fed6..ca83d50e 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -66,7 +66,7 @@ def common_options(parser, opts): opts.make_parser(parser, "scripts", metavar="SCRIPT", short="s") opts.make_parser(parser, "stickycookie", metavar="FILTER") opts.make_parser(parser, "stickyauth", metavar="FILTER") - opts.make_parser(parser, "streamfile", metavar="PATH", short="w") + opts.make_parser(parser, "save_stream_file", metavar="PATH", short="w") opts.make_parser(parser, "anticomp") # Proxy options @@ -128,7 +128,7 @@ def mitmdump(opts): nargs="...", help=""" Filter expression, equivalent to setting both the view_filter - and streamfile_filter options. + and save_stream_filter options. """ ) return parser diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py index f1170f9f..00e5cf4e 100644 --- a/mitmproxy/tools/console/flowlist.py +++ b/mitmproxy/tools/console/flowlist.py @@ -411,13 +411,13 @@ class FlowListBox(urwid.ListBox): val = not self.master.options.console_order_reversed self.master.options.console_order_reversed = val elif key == "W": - if self.master.options.streamfile: - self.master.options.streamfile = None + if self.master.options.save_stream_file: + self.master.options.save_stream_file = None else: signals.status_prompt_path.send( self, prompt="Stream flows to", - callback= lambda path: self.master.options.update(streamfile=path) + callback= lambda path: self.master.options.update(save_stream_file=path) ) else: return urwid.ListBox.keypress(self, size, key) diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py index d2424832..1930fa2f 100644 --- a/mitmproxy/tools/console/statusbar.py +++ b/mitmproxy/tools/console/statusbar.py @@ -251,8 +251,8 @@ class StatusBar(urwid.WidgetWrap): r.append(("heading_key", "s")) r.append("cripts:%s]" % len(self.master.options.scripts)) - if self.master.options.streamfile: - r.append("[W:%s]" % self.master.options.streamfile) + if self.master.options.save_stream_file: + r.append("[W:%s]" % self.master.options.save_stream_file) return r diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 9621dbbc..fefdca5c 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -127,7 +127,7 @@ def mitmdump(args=None): # pragma: no cover v = " ".join(args.filter_args) return dict( view_filter = v, - streamfile_filter = v, + save_stream_filter = v, ) return {} diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py new file mode 100644 index 00000000..fcd2feab --- /dev/null +++ b/test/mitmproxy/addons/test_save.py @@ -0,0 +1,62 @@ +import pytest + +from mitmproxy.test import taddons +from mitmproxy.test import tflow + +from mitmproxy import io +from mitmproxy import exceptions +from mitmproxy import options +from mitmproxy.addons import save + + +def test_configure(tmpdir): + sa = save.Save() + with taddons.context(options=options.Options()) as tctx: + with pytest.raises(exceptions.OptionsError): + tctx.configure(sa, save_stream_file=str(tmpdir)) + with pytest.raises(Exception, match="Invalid filter"): + tctx.configure( + sa, save_stream_file=str(tmpdir.join("foo")), save_stream_filter="~~" + ) + tctx.configure(sa, save_stream_filter="foo") + assert sa.filt + tctx.configure(sa, save_stream_filter=None) + assert not sa.filt + + +def rd(p): + x = io.FlowReader(open(p, "rb")) + return list(x.stream()) + + +def test_tcp(tmpdir): + sa = save.Save() + with taddons.context() as tctx: + p = str(tmpdir.join("foo")) + tctx.configure(sa, save_stream_file=p) + + tt = tflow.ttcpflow() + sa.tcp_start(tt) + sa.tcp_end(tt) + tctx.configure(sa, save_stream_file=None) + assert rd(p) + + +def test_simple(tmpdir): + sa = save.Save() + with taddons.context() as tctx: + p = str(tmpdir.join("foo")) + + tctx.configure(sa, save_stream_file=p) + + f = tflow.tflow(resp=True) + sa.request(f) + sa.response(f) + tctx.configure(sa, save_stream_file=None) + assert rd(p)[0].response + + tctx.configure(sa, save_stream_file="+" + p) + f = tflow.tflow() + sa.request(f) + tctx.configure(sa, save_stream_file=None) + assert not rd(p)[1].response diff --git a/test/mitmproxy/addons/test_streamfile.py b/test/mitmproxy/addons/test_streamfile.py deleted file mode 100644 index bcb27c79..00000000 --- a/test/mitmproxy/addons/test_streamfile.py +++ /dev/null @@ -1,62 +0,0 @@ -import pytest - -from mitmproxy.test import taddons -from mitmproxy.test import tflow - -from mitmproxy import io -from mitmproxy import exceptions -from mitmproxy import options -from mitmproxy.addons import streamfile - - -def test_configure(tmpdir): - sa = streamfile.StreamFile() - with taddons.context(options=options.Options()) as tctx: - with pytest.raises(exceptions.OptionsError): - tctx.configure(sa, streamfile=str(tmpdir)) - with pytest.raises(Exception, match="Invalid filter"): - tctx.configure( - sa, streamfile=str(tmpdir.join("foo")), streamfile_filter="~~" - ) - tctx.configure(sa, streamfile_filter="foo") - assert sa.filt - tctx.configure(sa, streamfile_filter=None) - assert not sa.filt - - -def rd(p): - x = io.FlowReader(open(p, "rb")) - return list(x.stream()) - - -def test_tcp(tmpdir): - sa = streamfile.StreamFile() - with taddons.context() as tctx: - p = str(tmpdir.join("foo")) - tctx.configure(sa, streamfile=p) - - tt = tflow.ttcpflow() - sa.tcp_start(tt) - sa.tcp_end(tt) - tctx.configure(sa, streamfile=None) - assert rd(p) - - -def test_simple(tmpdir): - sa = streamfile.StreamFile() - with taddons.context() as tctx: - p = str(tmpdir.join("foo")) - - tctx.configure(sa, streamfile=p) - - f = tflow.tflow(resp=True) - sa.request(f) - sa.response(f) - tctx.configure(sa, streamfile=None) - assert rd(p)[0].response - - tctx.configure(sa, streamfile="+" + p) - f = tflow.tflow() - sa.request(f) - tctx.configure(sa, streamfile=None) - assert not rd(p)[1].response -- cgit v1.2.3