From 9f77c79227f424f48ec16871d7efd451c8de6d47 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 3 Nov 2016 15:00:18 +1300 Subject: FileStreamer -> StreamFile options.wfile -> options.streamfile --- mitmproxy/addons/__init__.py | 4 +- mitmproxy/addons/filestreamer.py | 65 ---------------------------- mitmproxy/addons/streamfile.py | 65 ++++++++++++++++++++++++++++ mitmproxy/options.py | 4 +- mitmproxy/tools/cmdline.py | 16 +++---- mitmproxy/tools/console/flowlist.py | 6 +-- mitmproxy/tools/console/statusbar.py | 4 +- mitmproxy/tools/web/master.py | 6 +-- test/mitmproxy/addons/test_filestreamer.py | 69 ------------------------------ test/mitmproxy/addons/test_streamfile.py | 68 +++++++++++++++++++++++++++++ 10 files changed, 153 insertions(+), 154 deletions(-) delete mode 100644 mitmproxy/addons/filestreamer.py create mode 100644 mitmproxy/addons/streamfile.py delete mode 100644 test/mitmproxy/addons/test_filestreamer.py create mode 100644 test/mitmproxy/addons/test_streamfile.py diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index c5d40525..d2b50c35 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -1,7 +1,7 @@ from mitmproxy.addons import anticache from mitmproxy.addons import anticomp from mitmproxy.addons import clientplayback -from mitmproxy.addons import filestreamer +from mitmproxy.addons import streamfile from mitmproxy.addons import onboarding from mitmproxy.addons import replace from mitmproxy.addons import script @@ -20,7 +20,7 @@ def default_addons(): stickyauth.StickyAuth(), stickycookie.StickyCookie(), script.ScriptLoader(), - filestreamer.FileStreamer(), + streamfile.StreamFile(), streambodies.StreamBodies(), replace.Replace(), setheaders.SetHeaders(), diff --git a/mitmproxy/addons/filestreamer.py b/mitmproxy/addons/filestreamer.py deleted file mode 100644 index cb6f3762..00000000 --- a/mitmproxy/addons/filestreamer.py +++ /dev/null @@ -1,65 +0,0 @@ -import os.path - -from mitmproxy import exceptions -from mitmproxy import flowfilter -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): - 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, 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: - path, mode = options.outfile - if mode not in ("wb", "ab"): - raise exceptions.OptionsError("Invalid mode.") - 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 new file mode 100644 index 00000000..76befeb2 --- /dev/null +++ b/mitmproxy/addons/streamfile.py @@ -0,0 +1,65 @@ +import os.path + +from mitmproxy import exceptions +from mitmproxy import flowfilter +from mitmproxy import io + + +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, 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 "streamfile" in updated: + if self.stream: + self.done() + if options.streamfile: + path, mode = options.streamfile + if mode not in ("wb", "ab"): + raise exceptions.OptionsError("Invalid mode.") + 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 1db9f0f0..84aa4a74 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -48,7 +48,7 @@ class Options(optmanager.OptManager): stream_large_bodies: Optional[int] = None, verbosity: int = 2, default_contentview: str = "auto", - outfile: Optional[Tuple[str, str]] = None, + streamfile: Optional[Tuple[str, str]] = None, server_replay_ignore_content: bool = False, server_replay_ignore_params: Sequence[str] = (), server_replay_ignore_payload_params: Sequence[str] = (), @@ -108,7 +108,7 @@ class Options(optmanager.OptManager): self.stream_large_bodies = stream_large_bodies self.verbosity = verbosity self.default_contentview = default_contentview - self.outfile = outfile + self.streamfile = streamfile self.server_replay_ignore_content = server_replay_ignore_content self.server_replay_ignore_params = server_replay_ignore_params self.server_replay_ignore_payload_params = server_replay_ignore_payload_params diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index e4b29d0f..cb26378d 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -140,8 +140,8 @@ def get_common_options(args): raise exceptions.OptionsError(e) setheaders.append(p) - if args.outfile and args.outfile[0] == args.rfile: - if args.outfile[1] == "wb": + if args.streamfile and args.streamfile[0] == args.rfile: + if args.streamfile[1] == "wb": raise exceptions.OptionsError( "Cannot use '{}' for both reading and writing flows. " "Are you looking for --afile?".format(args.rfile) @@ -228,7 +228,7 @@ def get_common_options(args): stickyauth=stickyauth, stream_large_bodies=stream_large_bodies, showhost=args.showhost, - outfile=args.outfile, + streamfile=args.streamfile, verbosity=args.verbose, server_replay_nopop=args.server_replay_nopop, server_replay_ignore_content=args.server_replay_ignore_content, @@ -339,15 +339,15 @@ def basic_options(parser): action="store_const", dest="verbose", default=2, const=3, help="Increase log verbosity." ) - outfile = parser.add_mutually_exclusive_group() - outfile.add_argument( + streamfile = parser.add_mutually_exclusive_group() + streamfile.add_argument( "-w", "--wfile", - action="store", dest="outfile", type=lambda f: (f, "wb"), + action="store", dest="streamfile", type=lambda f: (f, "wb"), help="Write flows to file." ) - outfile.add_argument( + streamfile.add_argument( "-a", "--afile", - action="store", dest="outfile", type=lambda f: (f, "ab"), + action="store", dest="streamfile", type=lambda f: (f, "ab"), help="Append flows to file." ) parser.add_argument( diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py index 76545893..fc8368e7 100644 --- a/mitmproxy/tools/console/flowlist.py +++ b/mitmproxy/tools/console/flowlist.py @@ -393,13 +393,13 @@ class FlowListBox(urwid.ListBox): val = not self.master.options.order_reversed self.master.options.order_reversed = val elif key == "W": - if self.master.options.outfile: - self.master.options.outfile = None + if self.master.options.streamfile: + self.master.options.streamfile = None else: signals.status_prompt_path.send( self, prompt="Stream flows to", - callback= lambda path: self.master.options.update(outfile=(path, "ab")) + callback= lambda path: self.master.options.update(streamfile=(path, "ab")) ) else: return urwid.ListBox.keypress(self, size, key) diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py index e292cbd7..ebf2bca8 100644 --- a/mitmproxy/tools/console/statusbar.py +++ b/mitmproxy/tools/console/statusbar.py @@ -223,8 +223,8 @@ class StatusBar(urwid.WidgetWrap): r.append(("heading_key", "s")) r.append("cripts:%s]" % len(self.master.options.scripts)) - if self.master.options.outfile: - r.append("[W:%s]" % self.master.options.outfile[0]) + if self.master.options.streamfile: + r.append("[W:%s]" % self.master.options.streamfile[0]) return r diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index d283e5d5..3e1b5933 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -111,10 +111,10 @@ class WebMaster(master.Master): "error" ) - if options.outfile: + if options.streamfile: err = self.start_stream_to_path( - options.outfile[0], - options.outfile[1] + options.streamfile[0], + options.streamfile[1] ) if err: print("Stream file error: {}".format(err), file=sys.stderr) diff --git a/test/mitmproxy/addons/test_filestreamer.py b/test/mitmproxy/addons/test_filestreamer.py deleted file mode 100644 index 658a0aa8..00000000 --- a/test/mitmproxy/addons/test_filestreamer.py +++ /dev/null @@ -1,69 +0,0 @@ -from mitmproxy.test import tflow -from mitmproxy.test import tutils -from mitmproxy.test import taddons - -import os.path -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_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 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")) - - f = tflow.tflow(resp=True) - sa.request(f) - sa.response(f) - tctx.configure(sa, outfile=None) - assert rd(p)[0].response - - tctx.configure(sa, outfile=(p, "ab")) - f = tflow.tflow() - sa.request(f) - tctx.configure(sa, outfile=None) - assert not rd(p)[1].response - diff --git a/test/mitmproxy/addons/test_streamfile.py b/test/mitmproxy/addons/test_streamfile.py new file mode 100644 index 00000000..f63f29c4 --- /dev/null +++ b/test/mitmproxy/addons/test_streamfile.py @@ -0,0 +1,68 @@ +from mitmproxy.test import tflow +from mitmproxy.test import tutils +from mitmproxy.test import taddons + +import os.path +from mitmproxy import io +from mitmproxy import exceptions +from mitmproxy.tools import dump +from mitmproxy.addons import streamfile + + +def test_configure(): + sa = streamfile.StreamFile() + 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, streamfile=(tdir, "ab") + ) + tutils.raises( + "invalid filter", + tctx.configure, sa, streamfile=(p, "ab"), filtstr="~~" + ) + tutils.raises( + "invalid mode", + tctx.configure, sa, streamfile=(p, "xx") + ) + + +def rd(p): + x = io.FlowReader(open(p, "rb")) + return list(x.stream()) + + +def test_tcp(): + sa = streamfile.StreamFile() + with taddons.context() as tctx: + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") + tctx.configure(sa, streamfile=(p, "wb")) + + tt = tflow.ttcpflow() + sa.tcp_start(tt) + sa.tcp_end(tt) + tctx.configure(sa, streamfile=None) + assert rd(p) + + +def test_simple(): + sa = streamfile.StreamFile() + with taddons.context() as tctx: + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") + + tctx.configure(sa, streamfile=(p, "wb")) + + 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, "ab")) + f = tflow.tflow() + sa.request(f) + tctx.configure(sa, streamfile=None) + assert not rd(p)[1].response -- cgit v1.2.3