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: 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 ++++---
 2 files changed, 16 insertions(+), 6 deletions(-)

(limited to 'libmproxy')

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index d43da44b..7f6f1e7c 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 4025a30d..996b45cd 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):
         """
-- 
cgit v1.2.3