Massage namespace to make room for client replay.

Mostly replay -> server_replay

5 files changed, 45 insertions, 24 deletions

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 50d4dc4b..54be4933 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -9,7 +9,7 @@ class Options(object):
         "kill",
         "request_script",
         "response_script",
-        "replay",
+        "server_replay",
         "verbosity",
         "wfile",
         "rheaders",
@@ -64,14 +64,14 @@ class DumpMaster(flow.FlowMaster):
             except IOError, v:
                 raise DumpError(v.strerror)
 
-        if options.replay:
-            path = os.path.expanduser(options.replay)
+        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_playback(flows, options.kill, options.rheaders)
+            self.start_server_playback(flows, options.kill, options.rheaders)
 
 
     def _runscript(self, f, script):
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index a4e0a79c..2e734500 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -409,7 +409,8 @@ class FlowMaster(controller.Master):
     def __init__(self, server, state):
         controller.Master.__init__(self, server)
         self.state = state
-        self.playback = None
+        self.server_playback = None
+        self.client_playback = None
         self.scripts = {}
         self.kill_nonreplay = False
         self.stickycookie_state = False
@@ -434,21 +435,27 @@ class FlowMaster(controller.Master):
         else:
             self.stickycookie_state = None
 
-    def start_playback(self, flows, kill, headers):
+    def start_client_playback(self, flows):
+        """
+            flows: A list of flows.
+        """
+        self.client_playback = ClientPlaybackState(flows)
+
+    def start_server_playback(self, flows, kill, headers):
         """
             flows: A list of flows.
             kill: Boolean, should we kill requests not part of the replay?
         """
-        self.playback = ServerPlaybackState(headers, flows)
+        self.server_playback = ServerPlaybackState(headers, flows)
         self.kill_nonreplay = kill
 
-    def do_playback(self, flow):
+    def do_server_playback(self, flow):
         """
             This method should be called by child classes in the handle_request
             handler. Returns True if playback has taken place, None if not.
         """
-        if self.playback:
-            rflow = self.playback.next_flow(flow)
+        if self.server_playback:
+            rflow = self.server_playback.next_flow(flow)
             if not rflow:
                 return None
             response = proxy.Response.from_state(flow.request, rflow.response.get_state())
@@ -458,6 +465,11 @@ class FlowMaster(controller.Master):
             return True
         return None
 
+    def tick(self, q):
+        if self.client_playback:
+            self.client_playback.tick()
+        controller.Master.tick(self, q)
+
     def handle_clientconnect(self, r):
         self.state.clientconnect(r)
         r.ack()
@@ -468,6 +480,8 @@ class FlowMaster(controller.Master):
 
     def handle_error(self, r):
         f = self.state.add_error(r)
+        if self.client_playback:
+            self.client_playback.clear(f)
         r.ack()
         return f
 
@@ -477,8 +491,8 @@ class FlowMaster(controller.Master):
             self.stickycookie_state.handle_request(f)
         if "request" in self.scripts:
             self._runscript(f, self.scripts["request"])
-        if self.playback:
-            pb = self.do_playback(f)
+        if self.server_playback:
+            pb = self.do_server_playback(f)
             if not pb:
                 if self.kill_nonreplay:
                     self.state.kill_flow(f)
@@ -488,6 +502,8 @@ class FlowMaster(controller.Master):
 
     def handle_response(self, r):
         f = self.state.add_response(r)
+        if self.client_playback:
+            self.client_playback.clear(f)
         if not f:
             r.ack()
         if self.stickycookie_state:
diff --git a/mitmdump b/mitmdump
index f92444b8..63d942af 100755
--- a/mitmdump
+++ b/mitmdump
@@ -58,7 +58,7 @@ if __name__ == '__main__':
 
 
     group = OptionGroup(parser, "Server Replay")
-    group.add_option("-r", action="store", dest="replay", default=None, metavar="PATH",
+    group.add_option("-r", action="store", dest="server_replay", default=None, metavar="PATH",
                       help="Replay server responses from a saved file.")
     group.add_option("-k", "--kill",
                       action="store_true", dest="kill", default=False,
@@ -89,7 +89,7 @@ if __name__ == '__main__':
         wfile = options.wfile,
         request_script = options.request_script,
         response_script = options.response_script,
-        replay = options.replay,
+        server_replay = options.server_replay,
         kill = options.kill,
         rheaders = options.rheaders,
         stickycookie = stickycookie
diff --git a/test/test_dump.py b/test/test_dump.py
index cfa690b1..0a42c8c3 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -36,7 +36,7 @@ class uDumpMaster(libpry.AutoTree):
     def test_replay(self):
         cs = StringIO()
 
-        o = dump.Options(replay="nonexistent", kill=True)
+        o = dump.Options(server_replay="nonexistent", kill=True)
         libpry.raises(dump.DumpError, dump.DumpMaster, None, o, None, outfile=cs)
 
         t = self.tmpdir()
@@ -48,13 +48,13 @@ class uDumpMaster(libpry.AutoTree):
         fw.add(t)
         f.close()
 
-        o = dump.Options(replay=p, kill=True)
+        o = dump.Options(server_replay=p, kill=True)
         m = dump.DumpMaster(None, o, None, outfile=cs)
         
         self._cycle(m, "content")
         self._cycle(m, "content")
 
-        o = dump.Options(replay=p, kill=False)
+        o = dump.Options(server_replay=p, kill=False)
         m = dump.DumpMaster(None, o, None, outfile=cs)
         self._cycle(m, "nonexistent")
 
diff --git a/test/test_flow.py b/test/test_flow.py
index f9fb58e8..bd8c01a3 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -375,7 +375,7 @@ class uFlowMaster(libpry.AutoTree):
         err = proxy.Error(f.request, "msg")
         fm.handle_error(err)
 
-    def test_replay(self):
+    def test_server_playback(self):
         s = flow.State()
 
         f = utils.tflow()
@@ -383,15 +383,21 @@ class uFlowMaster(libpry.AutoTree):
         pb = [f]
 
         fm = flow.FlowMaster(None, s)
-        assert not fm.do_playback(utils.tflow())
+        assert not fm.do_server_playback(utils.tflow())
 
-        fm.start_playback(pb, False, [])
-        assert fm.do_playback(utils.tflow())
+        fm.start_server_playback(pb, False, [])
+        assert fm.do_server_playback(utils.tflow())
 
-        fm.start_playback(pb, False, [])
+        fm.start_server_playback(pb, False, [])
         r = utils.tflow()
         r.request.content = "gibble"
-        assert not fm.do_playback(r)
+        assert not fm.do_server_playback(r)
+
+    def test_client_playback(self):
+        s = flow.State()
+        fm = flow.FlowMaster(None, s)
+        pb = [utils.tflow_full()]
+        fm.start_client_playback(pb)
 
     def test_stickycookie(self):
         s = flow.State()
@@ -424,5 +430,4 @@ tests = [
     uState(),
     uSerialize(),
     uFlowMaster()
-
 ]