From 9bf1d300d8f4f1c258df5c1bb308546a81ee4df0 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sun, 17 Jul 2016 16:07:29 +1200 Subject: Some work on stream-to-file - stream -> filestreamer throughout - move active_flows set into FileStreamer addon --- mitmproxy/builtins/__init__.py | 4 +- mitmproxy/builtins/filestreamer.py | 66 ++++++++++++++++++++++++++++ mitmproxy/builtins/stream.py | 54 ----------------------- mitmproxy/flow/master.py | 7 +-- test/mitmproxy/builtins/test_filestreamer.py | 46 +++++++++++++++++++ test/mitmproxy/builtins/test_stream.py | 46 ------------------- 6 files changed, 116 insertions(+), 107 deletions(-) create mode 100644 mitmproxy/builtins/filestreamer.py delete mode 100644 mitmproxy/builtins/stream.py create mode 100644 test/mitmproxy/builtins/test_filestreamer.py delete mode 100644 test/mitmproxy/builtins/test_stream.py diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py index b4d3c0ff..1f8ae862 100644 --- a/mitmproxy/builtins/__init__.py +++ b/mitmproxy/builtins/__init__.py @@ -2,10 +2,10 @@ from __future__ import absolute_import, print_function, division from mitmproxy.builtins import anticache from mitmproxy.builtins import anticomp +from mitmproxy.builtins import filestreamer from mitmproxy.builtins import stickyauth from mitmproxy.builtins import stickycookie from mitmproxy.builtins import script -from mitmproxy.builtins import stream from mitmproxy.builtins import replace @@ -16,6 +16,6 @@ def default_addons(): stickyauth.StickyAuth(), stickycookie.StickyCookie(), script.ScriptLoader(), - stream.Stream(), + filestreamer.FileStreamer(), replace.Replace(), ] diff --git a/mitmproxy/builtins/filestreamer.py b/mitmproxy/builtins/filestreamer.py new file mode 100644 index 00000000..97ddc7c4 --- /dev/null +++ b/mitmproxy/builtins/filestreamer.py @@ -0,0 +1,66 @@ +from __future__ import absolute_import, print_function, division +import os.path + +from mitmproxy import exceptions +from mitmproxy.flow import io + + +class FileStreamer: + def __init__(self): + self.stream = None + self.active_flows = set() # type: Set[models.Flow] + + def start_stream_to_path(self, path, mode, filt): + path = os.path.expanduser(path) + try: + f = open(path, mode) + except IOError as v: + return str(v) + self.stream = io.FilteredFlowWriter(f, filt) + self.active_flows = set() + + def configure(self, options): + # We're already streaming - stop the previous stream and restart + if self.stream: + self.done() + + if options.outfile: + filt = None + if options.get("filtstr"): + filt = filt.parse(options.filtstr) + if not filt: + 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, filt) + if err: + raise exceptions.OptionsError(err) + + def tcp_open(self, flow): + if self.stream: + self.active_flows.add(flow) + + def tcp_close(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/builtins/stream.py b/mitmproxy/builtins/stream.py deleted file mode 100644 index 821a71f1..00000000 --- a/mitmproxy/builtins/stream.py +++ /dev/null @@ -1,54 +0,0 @@ -from __future__ import absolute_import, print_function, division -import os.path - -from mitmproxy import ctx -from mitmproxy import exceptions -from mitmproxy.flow import io - - -class Stream: - def __init__(self): - self.stream = None - - def start_stream_to_path(self, path, mode, filt): - path = os.path.expanduser(path) - try: - f = open(path, mode) - except IOError as v: - return str(v) - self.stream = io.FilteredFlowWriter(f, filt) - - def configure(self, options): - # We're already streaming - stop the previous stream and restart - if self.stream: - self.done() - - if options.outfile: - filt = None - if options.get("filtstr"): - filt = filt.parse(options.filtstr) - if not filt: - 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, filt) - if err: - raise exceptions.OptionsError(err) - - def done(self): - if self.stream: - for flow in ctx.master.active_flows: - self.stream.add(flow) - self.stream.fo.close() - self.stream = None - - def tcp_close(self, flow): - if self.stream: - self.stream.add(flow) - - def response(self, flow): - if self.stream: - self.stream.add(flow) diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index 80949825..244f24f6 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -30,7 +30,6 @@ class FlowMaster(controller.Master): if server: self.add_server(server) self.state = state - self.active_flows = set() # type: Set[models.Flow] self.server_playback = None # type: Optional[modules.ServerPlaybackState] self.client_playback = None # type: Optional[modules.ClientPlaybackState] self.kill_nonreplay = False @@ -329,7 +328,6 @@ class FlowMaster(controller.Master): return if f not in self.state.flows: # don't add again on replay self.state.add_flow(f) - self.active_flows.add(f) if not f.reply.acked: self.setheaders.run(f) if not f.reply.acked: @@ -348,7 +346,6 @@ class FlowMaster(controller.Master): @controller.handler def response(self, f): - self.active_flows.discard(f) self.state.update_flow(f) if not f.reply.acked: self.setheaders.run(f) @@ -367,7 +364,7 @@ class FlowMaster(controller.Master): def tcp_open(self, flow): # TODO: This would break mitmproxy currently. # self.state.add_flow(flow) - self.active_flows.add(flow) + pass @controller.handler def tcp_message(self, flow): @@ -382,4 +379,4 @@ class FlowMaster(controller.Master): @controller.handler def tcp_close(self, flow): - self.active_flows.discard(flow) + pass diff --git a/test/mitmproxy/builtins/test_filestreamer.py b/test/mitmproxy/builtins/test_filestreamer.py new file mode 100644 index 00000000..002006b7 --- /dev/null +++ b/test/mitmproxy/builtins/test_filestreamer.py @@ -0,0 +1,46 @@ +from __future__ import absolute_import, print_function, division + +from .. import tutils, mastertest + +import os.path + +from mitmproxy.builtins import filestreamer +from mitmproxy.flow import master, FlowReader +from mitmproxy.flow import state +from mitmproxy.flow import options + + +class TestStream(mastertest.MasterTest): + def test_stream(self): + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") + + def r(): + r = FlowReader(open(p, "rb")) + return list(r.stream()) + + s = state.State() + m = master.FlowMaster( + options.Options( + outfile = (p, "wb") + ), + None, + s + ) + sa = filestreamer.FileStreamer() + + m.addons.add(sa) + f = tutils.tflow(resp=True) + self.invoke(m, "request", f) + self.invoke(m, "response", f) + m.addons.remove(sa) + + assert r()[0].response + + m.options.outfile = (p, "ab") + + m.addons.add(sa) + f = tutils.tflow() + self.invoke(m, "request", f) + m.addons.remove(sa) + assert not r()[1].response diff --git a/test/mitmproxy/builtins/test_stream.py b/test/mitmproxy/builtins/test_stream.py deleted file mode 100644 index edaa41d2..00000000 --- a/test/mitmproxy/builtins/test_stream.py +++ /dev/null @@ -1,46 +0,0 @@ -from __future__ import absolute_import, print_function, division - -from .. import tutils, mastertest - -import os.path - -from mitmproxy.builtins import stream -from mitmproxy.flow import master, FlowReader -from mitmproxy.flow import state -from mitmproxy.flow import options - - -class TestStream(mastertest.MasterTest): - def test_stream(self): - with tutils.tmpdir() as tdir: - p = os.path.join(tdir, "foo") - - def r(): - r = FlowReader(open(p, "rb")) - return list(r.stream()) - - s = state.State() - m = master.FlowMaster( - options.Options( - outfile = (p, "wb") - ), - None, - s - ) - sa = stream.Stream() - - m.addons.add(sa) - f = tutils.tflow(resp=True) - self.invoke(m, "request", f) - self.invoke(m, "response", f) - m.addons.remove(sa) - - assert r()[0].response - - m.options.outfile = (p, "ab") - - m.addons.add(sa) - f = tutils.tflow() - self.invoke(m, "request", f) - m.addons.remove(sa) - assert not r()[1].response -- cgit v1.2.3