From fe99871df873f755ef5f4770edf37304f88102d7 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Mon, 21 Feb 2011 11:08:35 +1300
Subject: [PATCH] Add  --kill option to mitmdump

If this option is passed all requests that are not part of a replayed
conversation are killed. If the option is not passed, such requests are passed
through to the server as usual.
---
 libmproxy/dump.py | 15 ++++++++++++---
 libmproxy/flow.py |  7 ++++---
 mitmdump          |  4 ++++
 test/test_dump.py | 14 +++++++++++---
 test/test_flow.py |  4 ++--
 test/utils.py     |  7 +++++++
 6 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index d43da44bc..7f6f1e7cb 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -6,11 +6,12 @@ class DumpError(Exception): pass
 
 class Options(object):
     __slots__ = [
-        "verbosity",
-        "wfile",
+        "kill",
         "request_script",
         "response_script",
         "replay",
+        "verbosity",
+        "wfile",
     ]
     def __init__(self, **kwargs):
         for k, v in kwargs.items():
@@ -66,7 +67,15 @@ class DumpMaster(flow.FlowMaster):
         f = flow.FlowMaster.handle_request(self, r)
         if self.o.request_script:
             self._runscript(f, self.o.request_script)
-        if not self.playback(f):
+
+        if self.o.replay:
+            pb = self.playback(f)
+            if not pb:
+                if self.o.kill:
+                    self.state.kill_flow(f)
+                else:
+                    r.ack()
+        else:
             r.ack()
 
     def indent(self, n, t):
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 4025a30d3..996b45cdc 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -44,9 +44,10 @@ class ServerPlaybackState:
             chronological order.
         """
         for i in flows:
-            h = self._hash(i)
-            l = self.fmap.setdefault(self._hash(i), [])
-            l.append(i)
+            if i.response:
+                h = self._hash(i)
+                l = self.fmap.setdefault(self._hash(i), [])
+                l.append(i)
 
     def _hash(self, flow):
         """
diff --git a/mitmdump b/mitmdump
index c6a9ce107..a1b8aef10 100755
--- a/mitmdump
+++ b/mitmdump
@@ -50,6 +50,9 @@ if __name__ == '__main__':
     parser.add_option("-r", "--replay",
                       action="store", dest="replay", default=None,
                       help="Replay server responses from a saved file.")
+    parser.add_option("-k", "--kill",
+                      action="store_true", dest="kill", default=False,
+                      help="Kill extra requests during replay.")
 
     options, args = parser.parse_args()
 
@@ -65,6 +68,7 @@ if __name__ == '__main__':
         request_script = options.request_script,
         response_script = options.response_script,
         replay = options.replay,
+        kill = options.kill
     )
     if args:
         filt = " ".join(args)
diff --git a/test/test_dump.py b/test/test_dump.py
index 905404085..6213f8708 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -26,19 +26,27 @@ class uDumpMaster(libpry.AutoTree):
     def test_replay(self):
         cs = StringIO()
 
-        o = dump.Options(replay="nonexistent")
+        o = dump.Options(replay="nonexistent", kill=True)
         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()
+        t = utils.tflow_full()
+        t.response = utils.tresp(t.request)
         fw.add(t)
         f.close()
 
-        o = dump.Options(replay=p)
+        o = dump.Options(replay=p, kill=True)
         m = dump.DumpMaster(None, o, None, outfile=cs)
+        
+        self._cycle(m, "content")
+        self._cycle(m, "content")
+
+        o = dump.Options(replay=p, kill=False)
+        m = dump.DumpMaster(None, o, None, outfile=cs)
+        self._cycle(m, "nonexistent")
 
     def test_options(self):
         o = dump.Options(verbosity = 2)
diff --git a/test/test_flow.py b/test/test_flow.py
index cd88464d6..4bfd5ad9d 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -19,10 +19,10 @@ class uServerPlaybackState(libpry.AutoTree):
 
     def test_load(self):
         s = flow.ServerPlaybackState()
-        r = utils.tflow()
+        r = utils.tflow_full()
         r.request.headers["key"] = ["one"]
 
-        r2 = utils.tflow()
+        r2 = utils.tflow_full()
         r2.request.headers["key"] = ["two"]
 
         s.load([r, r2])
diff --git a/test/utils.py b/test/utils.py
index 126461067..9a00983c9 100644
--- a/test/utils.py
+++ b/test/utils.py
@@ -22,3 +22,10 @@ def tflow():
     return flow.Flow(r)
 
 
+def tflow_full():
+    r = treq()
+    f = flow.Flow(r)
+    f.response = tresp(r)
+    return f
+
+