From fd4dd8cb6b9e4e2a0afe0ecbf1bff52c66ce4dba Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Mon, 21 Feb 2011 09:54:39 +1300
Subject: First pass of playback function for mitmdump.

---
 libmproxy/dump.py  | 13 ++++++++++++-
 libmproxy/flow.py  | 21 +++++++++++++++++++++
 libmproxy/proxy.py | 12 +++++++++++-
 3 files changed, 44 insertions(+), 2 deletions(-)

(limited to 'libmproxy')

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index f6a7ae7e..d43da44b 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -10,6 +10,7 @@ class Options(object):
         "wfile",
         "request_script",
         "response_script",
+        "replay",
     ]
     def __init__(self, **kwargs):
         for k, v in kwargs.items():
@@ -38,6 +39,15 @@ class DumpMaster(flow.FlowMaster):
             except IOError, v:
                 raise DumpError(v.strerror)
 
+        if options.replay:
+            path = os.path.expanduser(options.replay)
+            try:
+                f = file(path, "r")
+                flows = list(flow.FlowReader(f).stream())
+            except IOError, v:
+                raise DumpError(v.strerror)
+            self.start_playback(flows)
+
     def _runscript(self, f, script):
         try:
             ret = f.run_script(script)
@@ -56,7 +66,8 @@ class DumpMaster(flow.FlowMaster):
         f = flow.FlowMaster.handle_request(self, r)
         if self.o.request_script:
             self._runscript(f, self.o.request_script)
-        r.ack()
+        if not self.playback(f):
+            r.ack()
 
     def indent(self, n, t):
         l = str(t).strip().split("\n")
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 16a2714c..4025a30d 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -328,6 +328,27 @@ class FlowMaster(controller.Master):
     def __init__(self, server, state):
         controller.Master.__init__(self, server)
         self.state = state
+        self._playback_state = None
+
+    def start_playback(self, flows):
+        self._playback_state = ServerPlaybackState()
+        self._playback_state.load(flows)
+
+    def 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_state:
+            rflow = self._playback_state.next_flow(flow)
+            if not rflow:
+                return None
+            response = proxy.Response.from_state(flow.request, rflow.response.get_state())
+            response.set_replay()
+            flow.response = response
+            flow.request.ack(response)
+            return True
+        return None
 
     def handle_clientconnect(self, r):
         self.state.clientconnect(r)
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index ed1c3d60..5e698c5f 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -265,6 +265,13 @@ class Response(controller.Msg):
         self.timestamp = timestamp or time.time()
         self.cached = False
         controller.Msg.__init__(self)
+        self.replay = False
+
+    def set_replay(self):
+        self.replay = True
+
+    def is_replay(self):
+        return self.replay
 
     def load_state(self, state):
         self.code = state["code"]
@@ -308,7 +315,10 @@ class Response(controller.Msg):
         return self.cached
 
     def short(self):
-        return "%s %s"%(self.code, self.msg)
+        r = "%s %s"%(self.code, self.msg)
+        if self.is_replay():
+            r = "[replay] " + r
+        return r
 
     def assemble(self):
         """
-- 
cgit v1.2.3