From fb28e71f0b1beb68723a2d89e008793325cee715 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Tue, 15 Mar 2011 15:11:03 +1300
Subject: [PATCH] Do pre-processing of requests before replay.

This enables scripts, anticache, server playback and sticky cookies for
request replays.
---
 libmproxy/console.py |  4 +--
 libmproxy/flow.py    | 67 ++++++++++++++++++++++++--------------------
 test/test_console.py |  1 +
 3 files changed, 39 insertions(+), 33 deletions(-)

diff --git a/libmproxy/console.py b/libmproxy/console.py
index b7e8f8796..9dd80c7ea 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -163,7 +163,7 @@ class ConnectionItem(WWrap):
                 self.master.statusbar.message("Can't delete connection mid-intercept.")
             self.master.sync_list_view()
         elif key == "r":
-            r = self.state.replay_request(self.flow, self.master.masterq)
+            r = self.master.replay_request(self.flow)
             if r:
                 self.master.statusbar.message(r)
             self.master.sync_list_view()
@@ -478,7 +478,7 @@ class ConnectionView(WWrap):
         elif key == "p":
             self.master.view_prev_flow(self.flow)
         elif key == "r":
-            r = self.state.replay_request(self.flow, self.master.masterq)
+            r = self.master.replay_request(self.flow)
             if r:
                 self.master.statusbar.message(r)
             self.master.refresh_connection(self.flow)
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index fb5c82d7b..21d4017ad 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -62,7 +62,7 @@ class ClientPlaybackState:
             self.current = master.handle_request(n.request)
             if not testing and not self.current.response:
                 #begin nocover
-                master.state.replay_request(self.current, master.masterq)
+                master.replay_request(self.current)
                 #end nocover
             elif self.current.response:
                 master.handle_response(self.current.response)
@@ -421,23 +421,6 @@ class State:
     def revert(self, f):
         f.revert()
 
-    def replay_request(self, f, masterq):
-        """
-            Returns None if successful, or error message if not.
-        """
-        #begin nocover
-        if f.intercepting:
-            return "Can't replay while intercepting..."
-        if f.request:
-            f.backup()
-            f.request.set_replay()
-            if f.request.content:
-                f.request.headers["content-length"] = [str(len(f.request.content))]
-            f.response = None
-            f.error = None
-            rt = RequestReplayThread(f, masterq)
-            rt.start()
-        #end nocover
 
 
 class FlowMaster(controller.Master):
@@ -526,6 +509,40 @@ class FlowMaster(controller.Master):
 
         controller.Master.tick(self, q)
 
+    def process_new_request(self, f):
+        if self.stickycookie_state:
+            self.stickycookie_state.handle_request(f)
+        if "request" in self.scripts:
+            self._runscript(f, self.scripts["request"])
+        if self.anticache:
+            f.request.anticache()
+        if self.server_playback:
+            pb = self.do_server_playback(f)
+            if not pb:
+                if self.kill_nonreplay:
+                    f.kill(self)
+                else:
+                    f.request.ack()
+
+    def replay_request(self, f):
+        """
+            Returns None if successful, or error message if not.
+        """
+        #begin nocover
+        if f.intercepting:
+            return "Can't replay while intercepting..."
+        if f.request:
+            f.backup()
+            f.request.set_replay()
+            if f.request.content:
+                f.request.headers["content-length"] = [str(len(f.request.content))]
+            f.response = None
+            f.error = None
+            self.process_new_request(f)
+            rt = RequestReplayThread(f, self.masterq)
+            rt.start()
+        #end nocover
+
     def handle_clientconnect(self, r):
         self.state.clientconnect(r)
         r.ack()
@@ -543,19 +560,7 @@ class FlowMaster(controller.Master):
 
     def handle_request(self, r):
         f = self.state.add_request(r)
-        if self.stickycookie_state:
-            self.stickycookie_state.handle_request(f)
-        if "request" in self.scripts:
-            self._runscript(f, self.scripts["request"])
-        if self.anticache:
-            r.anticache()
-        if self.server_playback:
-            pb = self.do_server_playback(f)
-            if not pb:
-                if self.kill_nonreplay:
-                    f.kill(self)
-                else:
-                    r.ack()
+        self.process_new_request(f)
         return f
 
     def handle_response(self, r):
diff --git a/test/test_console.py b/test/test_console.py
index d2cfb214c..dfa809eda 100644
--- a/test/test_console.py
+++ b/test/test_console.py
@@ -102,6 +102,7 @@ class uformat_flow(libpry.AutoTree):
         f.backup()
 
         f.request.set_replay()
+        f.response.set_replay()
         assert ('method', '[replay]') in console.format_flow(f, True)
         assert ('method', '[replay]') in console.format_flow(f, True, True)