diff options
| -rw-r--r-- | mitmproxy/builtins/__init__.py | 4 | ||||
| -rw-r--r-- | mitmproxy/builtins/filestreamer.py (renamed from mitmproxy/builtins/stream.py) | 26 | ||||
| -rw-r--r-- | mitmproxy/flow/master.py | 7 | ||||
| -rw-r--r-- | test/mitmproxy/builtins/test_filestreamer.py (renamed from test/mitmproxy/builtins/test_stream.py) | 4 |
4 files changed, 25 insertions, 16 deletions
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/stream.py b/mitmproxy/builtins/filestreamer.py index 821a71f1..97ddc7c4 100644 --- a/mitmproxy/builtins/stream.py +++ b/mitmproxy/builtins/filestreamer.py @@ -1,14 +1,14 @@ 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: +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) @@ -17,6 +17,7 @@ class Stream: 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 @@ -38,17 +39,28 @@ class Stream: if err: raise exceptions.OptionsError(err) - def done(self): + def tcp_open(self, flow): if self.stream: - for flow in ctx.master.active_flows: - self.stream.add(flow) - self.stream.fo.close() - self.stream = None + 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/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_stream.py b/test/mitmproxy/builtins/test_filestreamer.py index edaa41d2..002006b7 100644 --- a/test/mitmproxy/builtins/test_stream.py +++ b/test/mitmproxy/builtins/test_filestreamer.py @@ -4,7 +4,7 @@ from .. import tutils, mastertest import os.path -from mitmproxy.builtins import stream +from mitmproxy.builtins import filestreamer from mitmproxy.flow import master, FlowReader from mitmproxy.flow import state from mitmproxy.flow import options @@ -27,7 +27,7 @@ class TestStream(mastertest.MasterTest): None, s ) - sa = stream.Stream() + sa = filestreamer.FileStreamer() m.addons.add(sa) f = tutils.tflow(resp=True) |