From 613e9a298ec2fc82a4811c35f806d129ff883f2a Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sun, 15 May 2011 11:22:35 +1200 Subject: Add a new flow loading mechanism. We now simulate the normal connection flow when we load flows. That means that we can run scripts, hooks, sticky cookies, etc. --- libmproxy/console.py | 5 ++--- libmproxy/flow.py | 25 +++++++++++++++++++++---- test/test_flow.py | 21 +++++++++++++++++++++ test/tutils.py | 7 +++++++ 4 files changed, 51 insertions(+), 7 deletions(-) diff --git a/libmproxy/console.py b/libmproxy/console.py index 63dbf4f0..9e065e4d 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -1127,11 +1127,10 @@ class ConsoleMaster(flow.FlowMaster): try: f = file(path, "r") fr = flow.FlowReader(f) - data = list(fr.stream()) - f.close() except IOError, v: return v.strerror - self.state.load_flows(data) + flow.FlowMaster.load_flows(self, fr) + f.close() if self.conn_list_view: self.sync_list_view() self.focus_current() diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 387c49f0..b53794c9 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -542,6 +542,20 @@ class FlowMaster(controller.Master): controller.Master.tick(self, q) + def load_flows(self, fr): + """ + Load flows from a FlowReader object. + """ + for i in fr.stream(): + if i.request: + f = self.state.add_request(i.request) + self.process_new_request(f) + if i.response: + f = self.state.add_response(i.response) + self.process_new_response(f) + if i.error: + f = self.state.add_error(i.error) + def process_new_request(self, f): if self.stickycookie_state: self.stickycookie_state.handle_request(f) @@ -560,6 +574,12 @@ class FlowMaster(controller.Master): else: f.request.ack() + def process_new_response(self, f): + if self.stickycookie_state: + self.stickycookie_state.handle_response(f) + if "response" in self.scripts: + self._runscript(f, self.scripts["response"]) + def replay_request(self, f): """ Returns None if successful, or error message if not. @@ -605,10 +625,7 @@ class FlowMaster(controller.Master): self.client_playback.clear(f) if not f: r.ack() - if self.stickycookie_state: - self.stickycookie_state.handle_response(f) - if "response" in self.scripts: - self._runscript(f, self.scripts["response"]) + self.process_new_response(f) return f diff --git a/test/test_flow.py b/test/test_flow.py index a6489462..b7a368d2 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -369,6 +369,19 @@ class uState(libpry.AutoTree): class uSerialize(libpry.AutoTree): + def _treader(self): + sio = StringIO() + w = flow.FlowWriter(sio) + for i in range(3): + f = tutils.tflow_full() + w.add(f) + for i in range(3): + f = tutils.tflow_err() + w.add(f) + + sio.seek(0) + return flow.FlowReader(sio) + def test_roundtrip(self): sio = StringIO() f = tutils.tflow() @@ -381,6 +394,14 @@ class uSerialize(libpry.AutoTree): assert len(l) == 1 assert l[0] == f + def test_load_flows(self): + r = self._treader() + s = flow.State() + fm = flow.FlowMaster(None, s) + fm.load_flows(r) + assert len(s.flow_list) == 6 + + def test_error(self): sio = StringIO() sio.write("bogus") diff --git a/test/tutils.py b/test/tutils.py index fea7f224..2327bdb9 100644 --- a/test/tutils.py +++ b/test/tutils.py @@ -32,6 +32,13 @@ def tflow_full(): return f +def tflow_err(): + r = treq() + f = flow.Flow(r) + f.error = proxy.Error(r, "error") + return f + + # Yes, the random ports are horrible. During development, sockets are often not # properly closed during error conditions, which means you have to wait until # you can re-bind to the same port. This is a pain in the ass, so we just pick -- cgit v1.2.3 From e285b17e3ff2e48455a5bc45b9a401b5b78b65cf Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sun, 15 May 2011 11:54:12 +1200 Subject: Add -r option to mitmdump and mitmproxy. This option reads a set of flows from a file. I've also regularized the mitmdump and mitmproxy command-line signatures by removing mitmproxy's old way of specifying flow loads through naked arguments. --- libmproxy/cmdline.py | 6 ++++++ libmproxy/console.py | 5 +++++ libmproxy/dump.py | 12 ++++++++++++ libmproxy/flow.py | 8 +++----- mitmdump | 2 +- mitmproxy | 9 +-------- 6 files changed, 28 insertions(+), 14 deletions(-) diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py index 222defeb..ce68baed 100644 --- a/libmproxy/cmdline.py +++ b/libmproxy/cmdline.py @@ -21,6 +21,7 @@ def get_common_options(options): no_server = options.no_server, refresh_server_playback = not options.norefresh, rheaders = options.rheaders, + rfile = options.rfile, request_script = options.request_script, response_script = options.response_script, server_replay = options.server_replay, @@ -57,6 +58,11 @@ def common_options(parser): action="store_true", dest="quiet", help="Quiet." ) + parser.add_option( + "-r", + action="store", dest="rfile", default=None, + help="Read flows from file." + ) parser.add_option( "--anticache", action="store_true", dest="anticache", default=False, diff --git a/libmproxy/console.py b/libmproxy/console.py index 9e065e4d..aa763930 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -782,6 +782,7 @@ class Options(object): "refresh_server_playback", "request_script", "response_script", + "rfile", "rheaders", "server_replay", "stickycookie", @@ -818,6 +819,7 @@ class ConsoleMaster(flow.FlowMaster): def __init__(self, server, options): flow.FlowMaster.__init__(self, server, ConsoleState()) self.looptime = 0 + self.options = options self.conn_list_view = None self.set_palette() @@ -1031,6 +1033,9 @@ class ConsoleMaster(flow.FlowMaster): slave = controller.Slave(self.masterq, self.server) slave.start() + if self.options.rfile: + self.load_flows(self.options.rfile) + self.ui.run_wrapper(self.loop) # If True, quit just pops out to connection list view. print >> sys.stderr, "Shutting down..." diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 3e645167..ce9c2b50 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -14,6 +14,7 @@ class Options(object): "refresh_server_playback", "request_script", "response_script", + "rfile", "rheaders", "server_replay", "stickycookie", @@ -84,6 +85,15 @@ class DumpMaster(flow.FlowMaster): not options.keepserving ) + if options.rfile: + path = os.path.expanduser(options.rfile) + try: + f = file(path, "r") + freader = flow.FlowReader(f) + except IOError, v: + raise DumpError(v.strerror) + self.load_flows(freader) + if options.client_replay: self.start_client_playback( self._readflow(options.client_replay), @@ -189,6 +199,8 @@ class DumpMaster(flow.FlowMaster): # begin nocover def run(self): + if self.o.rfile and not self.o.keepserving: + return try: return flow.FlowMaster.run(self) except BaseException, v: diff --git a/libmproxy/flow.py b/libmproxy/flow.py index b53794c9..14f06241 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -548,13 +548,11 @@ class FlowMaster(controller.Master): """ for i in fr.stream(): if i.request: - f = self.state.add_request(i.request) - self.process_new_request(f) + self.handle_request(i.request) if i.response: - f = self.state.add_response(i.response) - self.process_new_response(f) + self.handle_response(i.response) if i.error: - f = self.state.add_error(i.error) + self.handle_error(i.error) def process_new_request(self, f): if self.stickycookie_state: diff --git a/mitmdump b/mitmdump index b54d051f..bb4ec501 100755 --- a/mitmdump +++ b/mitmdump @@ -30,7 +30,7 @@ if __name__ == '__main__': parser.add_option( "--keepserving", action="store_true", dest="keepserving", default=False, - help="Continue serving after playback. We exit by default." + help="Continue serving after client playback or file read. We exit by default." ) options, args = parser.parse_args() diff --git a/mitmproxy b/mitmproxy index bcb079a3..942646b2 100755 --- a/mitmproxy +++ b/mitmproxy @@ -23,7 +23,7 @@ from optparse import OptionParser, OptionGroup if __name__ == '__main__': parser = OptionParser( - usage = "%prog [options] [flowdump path]", + usage = "%prog [options]", version="%%prog %s"%VERSION, ) cmdline.common_options(parser) @@ -65,13 +65,6 @@ if __name__ == '__main__': opts.debug = options.debug m = console.ConsoleMaster(server, opts) - for i in args: - try: - m.load_flows(i) - except flow.FlowReadError, v: - print >> sys.stderr, "mitmproxy:", v.args[0] - sys.exit(1) - m.run() -- cgit v1.2.3 From 07110bbbf192339b52d1067b531d0cea20b245d2 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sun, 15 May 2011 12:23:34 +1200 Subject: Anticache and refresh_server_playback options are applied before flows are loaded. You can now use mitmdump to preview how these options work, by running mitmdump against a set of saved flows, and viewing the output. --- libmproxy/dump.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libmproxy/dump.py b/libmproxy/dump.py index ce9c2b50..b4621716 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -53,6 +53,8 @@ class DumpMaster(flow.FlowMaster): flow.FlowMaster.__init__(self, server, flow.State()) self.outfile = outfile self.o = options + self.anticache = options.anticache + self.refresh_server_playback = options.refresh_server_playback if filtstr: self.filt = filt.parse(filtstr) @@ -100,8 +102,6 @@ class DumpMaster(flow.FlowMaster): not options.keepserving ) - self.anticache = options.anticache - self.refresh_server_playback = options.refresh_server_playback def _readflow(self, path): path = os.path.expanduser(path) -- cgit v1.2.3