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: [PATCH] First pass of playback function for mitmdump.

---
 libmproxy/dump.py  | 13 ++++++++++++-
 libmproxy/flow.py  | 21 +++++++++++++++++++++
 libmproxy/proxy.py | 12 +++++++++++-
 mitmdump           |  4 ++++
 test/test_dump.py  | 17 +++++++++++++++++
 test/test_flow.py  | 18 ++++++++++++++++++
 test/utils.py      |  2 ++
 7 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index f6a7ae7ed..d43da44bc 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 16a2714c6..4025a30d3 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 ed1c3d608..5e698c5fb 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):
         """
diff --git a/mitmdump b/mitmdump
index 477260915..c6a9ce107 100755
--- a/mitmdump
+++ b/mitmdump
@@ -47,6 +47,9 @@ if __name__ == '__main__':
     parser.add_option("", "--respscript",
                       action="store", dest="response_script", default=None,
                       help="Script to run when a response is recieved.")
+    parser.add_option("-r", "--replay",
+                      action="store", dest="replay", default=None,
+                      help="Replay server responses from a saved file.")
 
     options, args = parser.parse_args()
 
@@ -61,6 +64,7 @@ if __name__ == '__main__':
         wfile = options.wfile,
         request_script = options.request_script,
         response_script = options.response_script,
+        replay = options.replay,
     )
     if args:
         filt = " ".join(args)
diff --git a/test/test_dump.py b/test/test_dump.py
index 46dd6dfd2..905404085 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -23,6 +23,23 @@ class uDumpMaster(libpry.AutoTree):
         self._cycle(m, content)
         return cs.getvalue()
 
+    def test_replay(self):
+        cs = StringIO()
+
+        o = dump.Options(replay="nonexistent")
+        libpry.raises(dump.DumpError, dump.DumpMaster, None, o, None, outfile=cs)
+
+        t = self.tmpdir()
+        p = os.path.join(t, "rep")
+        f = open(p, "w")
+        fw = flow.FlowWriter(f)
+        t = utils.tflow()
+        fw.add(t)
+        f.close()
+
+        o = dump.Options(replay=p)
+        m = dump.DumpMaster(None, o, None, outfile=cs)
+
     def test_options(self):
         o = dump.Options(verbosity = 2)
         assert o.verbosity == 2
diff --git a/test/test_flow.py b/test/test_flow.py
index adfeda6ea..cd88464d6 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -309,6 +309,24 @@ class uFlowMaster(libpry.AutoTree):
         err = proxy.Error(f.request, "msg")
         fm.handle_error(err)
 
+    def test_replay(self):
+        s = flow.State()
+
+        f = utils.tflow()
+        f.response = utils.tresp(f.request)
+        pb = [f]
+
+        fm = flow.FlowMaster(None, s)
+        assert not fm.playback(utils.tflow())
+
+        fm.start_playback(pb)
+        assert fm.playback(utils.tflow())
+
+        fm.start_playback(pb)
+        r = utils.tflow()
+        r.request.content = "gibble"
+        assert not fm.playback(r)
+
 
 
 tests = [
diff --git a/test/utils.py b/test/utils.py
index b1dc46d49..126461067 100644
--- a/test/utils.py
+++ b/test/utils.py
@@ -1,3 +1,4 @@
+import os.path
 from libmproxy import proxy, utils, filt, flow
 
 def treq(conn=None):
@@ -20,3 +21,4 @@ def tflow():
     r = treq()
     return flow.Flow(r)
 
+