From 77a33c441b9d6ea80bda9a6288cfcba95ab70081 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sun, 19 Feb 2012 11:29:49 +1300
Subject: [PATCH] Add duplicate_flow and replay_request hooks to ScriptContext.

---
 libmproxy/controller.py | 11 ++++++-----
 libmproxy/flow.py       | 31 ++++++++++++++++++++++++++++---
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/libmproxy/controller.py b/libmproxy/controller.py
index 7607d6d6c..dac2c82aa 100644
--- a/libmproxy/controller.py
+++ b/libmproxy/controller.py
@@ -25,11 +25,12 @@ class Msg:
         self.acked = False
 
     def _ack(self, data=False):
-        self.acked = True
-        if data is None:
-            self.q.put(data)
-        else:
-            self.q.put(data or self)
+        if not self.acked:
+            self.acked = True
+            if data is None:
+                self.q.put(data)
+            else:
+                self.q.put(data or self)
 
     def _send(self, masterq):
         self.acked = False
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index db4cde5ee..66cb4790e 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -25,6 +25,24 @@ class ScriptContext:
         """
         self._master.add_event(*args, **kwargs)
 
+    def duplicate_flow(self, f):
+        """
+            Returns a duplicate of the specified flow. The flow is also
+            injected into the current state, and is ready for editing, replay,
+            etc.
+        """
+        self._master.pause_scripts = True
+        f = self._master.duplicate_flow(f)
+        self._master.pause_scripts = False
+        return f
+
+    def replay_request(self, f):
+        """
+            Replay the request on the current flow. The response will be added
+            to the flow object.
+        """
+        self._master.replay_request(f)
+
 
 class Headers:
     def __init__(self, lst=None):
@@ -300,6 +318,7 @@ class Request(HTTPMsg):
             Returns a copy of this object.
         """
         c = copy.copy(self)
+        c.acked = True
         c.headers = self.headers.copy()
         return c
 
@@ -521,6 +540,7 @@ class Response(HTTPMsg):
             Returns a copy of this object.
         """
         c = copy.copy(self)
+        c.acked = True
         c.headers = self.headers.copy()
         return c
 
@@ -613,7 +633,9 @@ class ClientConnect(controller.Msg):
         """
             Returns a copy of this object.
         """
-        return copy.copy(self)
+        c = copy.copy(self)
+        c.acked = True
+        return c
 
 
 class Error(controller.Msg):
@@ -644,7 +666,9 @@ class Error(controller.Msg):
         """
             Returns a copy of this object.
         """
-        return copy.copy(self)
+        c = copy.copy(self)
+        c.acked = True
+        return c
 
     def _get_state(self):
         return dict(
@@ -1125,6 +1149,7 @@ class FlowMaster(controller.Master):
         self.client_playback = None
         self.kill_nonreplay = False
         self.script = None
+        self.pause_scripts = False
 
         self.stickycookie_state = False
         self.stickycookie_txt = None
@@ -1318,7 +1343,7 @@ class FlowMaster(controller.Master):
         #end nocover
 
     def run_script_hook(self, name, *args, **kwargs):
-        if self.script:
+        if self.script and not self.pause_scripts:
             ret = self.script.run(name, *args, **kwargs)
             if not ret[0] and ret[1]:
                 e = "Script error:\n" + ret[1][1]