diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py index 727765428..4034c1889 100644 --- a/libmproxy/console/__init__.py +++ b/libmproxy/console/__init__.py @@ -928,7 +928,7 @@ class ConsoleMaster(flow.FlowMaster): self.state.killall(self) if self.stream: self.stream.fo.close() - controller.Master.shutdown(self) + flow.FlowMaster.shutdown(self) def sync_list_view(self): self.flow_list_walker._modified() diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 5f5cad4c0..17e88bc11 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -17,7 +17,7 @@ This module provides more sophisticated flow tracking. These match requests with their responses, and provide filtering and interception facilities. """ -import hashlib, Cookie, cookielib, copy, re, urlparse +import hashlib, Cookie, cookielib, copy, re, urlparse, os import time import tnetstring, filt, script, utils, encoding, proxy from email.utils import parsedate_tz, formatdate, mktime_tz @@ -1200,6 +1200,8 @@ class FlowMaster(controller.Master): self.refresh_server_playback = False self.replacehooks = ReplaceHooks() + self.stream = None + def add_event(self, e, level="info"): """ level: info, error @@ -1415,21 +1417,35 @@ class FlowMaster(controller.Master): def handle_response(self, r): f = self.state.add_response(r) - self.replacehooks.run(f) if f: + self.replacehooks.run(f) self.run_script_hook("response", f) - if self.client_playback: - self.client_playback.clear(f) - if not f: - r._ack() - if f: + if self.client_playback: + self.client_playback.clear(f) self.process_new_response(f) + if self.stream: + self.stream.add(f) + else: + r._ack() return f def shutdown(self): if self.script: self.load_script(None) controller.Master.shutdown(self) + if self.stream: + for i in self.state._flow_list: + if not i.response: + self.stream.add(i) + self.stream.fo.close() + self.stop_stream() + + def start_stream(self, fp): + self.stream = FlowWriter(fp) + + def stop_stream(self): + self.stream = None + class FlowWriter: diff --git a/test/test_flow.py b/test/test_flow.py index 170b16bb7..99df9ed0e 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -1,4 +1,4 @@ -import Queue, time +import Queue, time, os.path from cStringIO import StringIO import email.utils from libmproxy import filt, flow, controller, utils, tnetstring @@ -689,6 +689,32 @@ class TestFlowMaster: fm.handle_request(f.request) assert f.request.headers["authorization"] == ["foo"] + def test_stream(self): + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") + def r(): + r = flow.FlowReader(open(p)) + return list(r.stream()) + + s = flow.State() + fm = flow.FlowMaster(None, s) + tf = tutils.tflow_full() + + fm.start_stream(file(p, "ab")) + fm.handle_request(tf.request) + fm.handle_response(tf.response) + fm.stop_stream() + + assert r()[0].response + + tf = tutils.tflow_full() + fm.start_stream(file(p, "ab")) + fm.handle_request(tf.request) + fm.shutdown() + + assert not r()[1].response + + class TestRequest: def test_simple(self): h = flow.ODictCaseless()