From 5da4099ddfee61b717c23a8c4227947ecc2f6c53 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sat, 5 Mar 2011 13:03:26 +1300
Subject: Massage namespace to make room for client replay.

Mostly replay -> server_replay
---
 libmproxy/flow.py | 32 ++++++++++++++++++++++++--------
 1 file changed, 24 insertions(+), 8 deletions(-)

(limited to 'libmproxy/flow.py')

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:
-- 
cgit v1.2.3