From bad77f3470f11716f700db22f51b1f53285d23ca Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sun, 6 Mar 2011 11:21:31 +1300
Subject: Add client replay options to mitmdump.

---
 libmproxy/dump.py | 34 +++++++++++++++++++++++-----------
 libmproxy/flow.py |  7 +++----
 2 files changed, 26 insertions(+), 15 deletions(-)

(limited to 'libmproxy')

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 54be4933..621ffae5 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -10,6 +10,7 @@ class Options(object):
         "request_script",
         "response_script",
         "server_replay",
+        "client_replay",
         "verbosity",
         "wfile",
         "rheaders",
@@ -31,7 +32,11 @@ def str_response(resp):
 
 
 def str_request(req):
-    r = "%s %s"%(req.method, req.url())
+    if req.client_conn:
+        c = req.client_conn.address[0]
+    else:
+        c = "[replay]"
+    r = "%s %s %s"%(c, req.method, req.url())
     if req.stickycookie:
         r = "[stickycookie] " + r
     return r
@@ -65,14 +70,22 @@ class DumpMaster(flow.FlowMaster):
                 raise DumpError(v.strerror)
 
         if options.server_replay:
-            path = os.path.expanduser(options.server_replay)
-            try:
-                f = file(path, "r")
-                flows = list(flow.FlowReader(f).stream())
-            except IOError, v:
-                raise DumpError(v.strerror)
-            self.start_server_playback(flows, options.kill, options.rheaders)
+            self.start_server_playback(
+                self._readflow(options.server_replay),
+                options.kill, options.rheaders
+            )
 
+        if options.client_replay:
+            self.start_client_playback(self._readflow(options.client_replay))
+
+    def _readflow(self, path):
+        path = os.path.expanduser(path)
+        try:
+            f = file(path, "r")
+            flows = list(flow.FlowReader(f).stream())
+        except IOError, v:
+            raise DumpError(v.strerror)
+        return flows
 
     def _runscript(self, f, script):
         try:
@@ -92,6 +105,7 @@ class DumpMaster(flow.FlowMaster):
         f = flow.FlowMaster.handle_request(self, r)
         if f:
             r.ack()
+        return f
 
     def indent(self, n, t):
         l = str(t).strip().split("\n")
@@ -105,12 +119,10 @@ class DumpMaster(flow.FlowMaster):
                     return
             sz = utils.pretty_size(len(f.response.content))
             if self.o.verbosity == 1:
-                print >> self.outfile, f.request.client_conn.address[0],
                 print >> self.outfile, str_request(f.request)
                 print >> self.outfile, "  <<",
                 print >> self.outfile, str_response(f.response), sz
             elif self.o.verbosity == 2:
-                print >> self.outfile, f.request.client_conn.address[0],
                 print >> self.outfile, str_request(f.request)
                 print >> self.outfile, self.indent(4, f.request.headers)
                 print >> self.outfile
@@ -118,7 +130,6 @@ class DumpMaster(flow.FlowMaster):
                 print >> self.outfile, self.indent(4, f.response.headers)
                 print >> self.outfile, "\n"
             elif self.o.verbosity == 3:
-                print >> self.outfile, f.request.client_conn.address[0],
                 print >> self.outfile, str_request(f.request)
                 print >> self.outfile, self.indent(4, f.request.headers)
                 if utils.isBin(f.request.content):
@@ -136,6 +147,7 @@ class DumpMaster(flow.FlowMaster):
             self.state.delete_flow(f)
             if self.o.wfile:
                 self.fwriter.add(f)
+        return f
 
 # begin nocover
     def run(self):
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 2e734500..55976610 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -52,9 +52,8 @@ class ClientPlaybackState:
             testing: Disables actual replay for testing.
         """
         if self.flows and not self.current:
-            self.current = self.flows.pop(0)
-            self.current.response = None
-            master.handle_request(self.current.request)
+            n = self.flows.pop(0)
+            self.current = master.handle_request(n.request)
             if not testing:
                 #begin nocover
                 master.state.replay_request(self.current, master.masterq)
@@ -467,7 +466,7 @@ class FlowMaster(controller.Master):
 
     def tick(self, q):
         if self.client_playback:
-            self.client_playback.tick()
+            self.client_playback.tick(self)
         controller.Master.tick(self, q)
 
     def handle_clientconnect(self, r):
-- 
cgit v1.2.3