From 324d7c3955bef29745183949baed48ee286c126d Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sun, 13 Mar 2011 16:50:11 +1300
Subject: Add client plaback to mitmproxy.

---
 libmproxy/console.py | 36 ++++++++++++++++++++++++++++++++++--
 libmproxy/flow.py    |  6 +++---
 test/test_flow.py    |  6 ++++--
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/libmproxy/console.py b/libmproxy/console.py
index f13bc933..f068c163 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -675,6 +675,10 @@ class StatusBar(WWrap):
     def get_status(self):
         r = []
 
+        if self.master.client_playback:
+            r.append("[")
+            r.append(("statusbar_highlight", "cplayback"))
+            r.append(":%s to go]"%self.master.client_playback.count())
         if self.master.state.intercept_txt:
             r.append("[")
             r.append(("statusbar_highlight", "i"))
@@ -872,6 +876,24 @@ class ConsoleMaster(flow.FlowMaster):
         self.stickyhosts = {}
         self.anticache = options.anticache
 
+        if options.client_replay:
+            self.client_playback_path(options.client_replay)
+
+    def _readflow(self, path):
+        path = os.path.expanduser(path)
+        try:
+            f = file(path, "r")
+            flows = list(flow.FlowReader(f).stream())
+        except (IOError, flow.FlowReadError), v:
+            return True, v.strerror
+        return False, flows
+
+    def client_playback_path(self, path):
+        err, ret = self._readflow(path)
+        if err:
+            self.statusbar.message(ret)
+        else:
+            self.start_client_playback(ret, True)
 
     def spawn_external_viewer(self, data, contenttype):
         if contenttype:
@@ -1246,6 +1268,13 @@ class ConsoleMaster(flow.FlowMaster):
                         self.statusbar.message("")
                         if k == "?":
                             self.view_help()
+                        elif k == "c":
+                            self.path_prompt(
+                                "Client replay: ",
+                                self.state.last_saveload,
+                                self.client_playback_path
+                            )
+                            k = None
                         elif k == "l":
                             self.prompt("Limit: ", self.state.limit_txt, self.set_limit)
                             self.sync_list_view()
@@ -1329,7 +1358,7 @@ class ConsoleMaster(flow.FlowMaster):
 
     def shutdown(self):
         for i in self.state.flow_list:
-            i.kill()
+            i.kill(self)
         controller.Master.shutdown(self)
 
     def sync_list_view(self):
@@ -1344,7 +1373,7 @@ class ConsoleMaster(flow.FlowMaster):
         self.sync_list_view()
 
     def kill_connection(self, f):
-        self.state.kill_flow(f)
+        f.kill(self)
 
     def refresh_connection(self, c):
         if hasattr(self.header, "refresh_connection"):
@@ -1372,14 +1401,17 @@ class ConsoleMaster(flow.FlowMaster):
         f = flow.FlowMaster.handle_error(self, r)
         if f:
             self.process_flow(f, r)
+        return f
 
     def handle_request(self, r):
         f = flow.FlowMaster.handle_request(self, r)
         if f:
             self.process_flow(f, r)
+        return f
 
     def handle_response(self, r):
         f = flow.FlowMaster.handle_response(self, r)
         if f:
             self.process_flow(f, r)
+        return f
 
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 3764fc95..fb5c82d7 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -287,13 +287,14 @@ class Flow:
                 return pattern(self.request)
         return False
 
-    def kill(self):
+    def kill(self, master):
         self.error = proxy.Error(self.request, "Connection killed")
         if self.request and not self.request.acked:
             self.request.ack(None)
         elif self.response and not self.response.acked:
             self.response.ack(None)
         self.intercepting = False
+        master.handle_error(self.error)
 
     def intercept(self):
         self.intercepting = True
@@ -552,8 +553,7 @@ class FlowMaster(controller.Master):
             pb = self.do_server_playback(f)
             if not pb:
                 if self.kill_nonreplay:
-                    f.kill()
-                    self.handle_error(f.error)
+                    f.kill(self)
                 else:
                     r.ack()
         return f
diff --git a/test/test_flow.py b/test/test_flow.py
index 947d0d59..91e3f5f4 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -168,18 +168,20 @@ class uFlow(libpry.AutoTree):
         assert f == f2
 
     def test_kill(self):
+        s = flow.State()
+        fm = flow.FlowMaster(None, s)
         f = tutils.tflow()
         f.request = tutils.treq()
         f.intercept()
         assert not f.request.acked
-        f.kill()
+        f.kill(fm)
         assert f.request.acked
         f.intercept()
         f.response = tutils.tresp()
         f.request = f.response.request
         f.request.ack()
         assert not f.response.acked
-        f.kill()
+        f.kill(fm)
         assert f.response.acked
 
     def test_accept_intercept(self):
-- 
cgit v1.2.3