From f028dc7e1bb030035ecb72fa0e9f175910d83656 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Thu, 17 Feb 2011 10:44:08 +1300
Subject: Add filtering to mitmdump.

---
 libmproxy/dump.py | 15 ++++++++++-----
 mitmdump          |  9 +++++++--
 test/test_dump.py | 17 +++++++++++++----
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index e30c2422..ece6e190 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -1,5 +1,5 @@
 import sys, os
-import flow
+import flow, filt
 
 class DumpError(Exception): pass
 
@@ -18,11 +18,16 @@ class Options(object):
 
 
 class DumpMaster(flow.FlowMaster):
-    def __init__(self, server, options, outfile=sys.stdout):
+    def __init__(self, server, options, filtstr, outfile=sys.stdout):
         flow.FlowMaster.__init__(self, server, flow.State())
         self.outfile = outfile
         self.o = options
 
+        if filtstr:
+            self.filt = filt.parse(filtstr)
+        else:
+            self.filt = None
+
         if options.wfile:
             path = os.path.expanduser(options.wfile)
             try:
@@ -46,6 +51,9 @@ class DumpMaster(flow.FlowMaster):
     def handle_response(self, msg):
         f = flow.FlowMaster.handle_response(self, msg)
         if f:
+            msg.ack()
+            if self.filt and not f.match(self.filt):
+                    return
             if 0 < self.o.verbosity < 3:
                 print >> self.outfile, ">>",
                 print >> self.outfile, msg.request.short()
@@ -66,10 +74,7 @@ class DumpMaster(flow.FlowMaster):
                 for i in msg.assemble().splitlines():
                     print >> self.outfile, "\t", i
                 print >> self.outfile, "<<"
-
-            msg.ack()
             self.state.delete_flow(f)
-
             if self.o.wfile:
                 self.fwriter.add(f)
 
diff --git a/mitmdump b/mitmdump
index 47fe7c2d..886bc5cc 100755
--- a/mitmdump
+++ b/mitmdump
@@ -23,7 +23,7 @@ from optparse import OptionParser, OptionGroup
 
 if __name__ == '__main__':
     parser = OptionParser(
-                usage = "%prog [options]",
+                usage = "%prog [options] [filter]",
                 version="%%prog %s"%VERSION,
             )
     utils.certificate_option_group(parser)
@@ -60,7 +60,12 @@ if __name__ == '__main__':
         verbosity = options.verbose,
         wfile = options.wfile
     )
-    m = dump.DumpMaster(server, dumpopts)
+    if args:
+        filt = " ".join(args)
+    else:
+        filt = None
+
+    m = dump.DumpMaster(server, dumpopts, filt)
 
     if options.verbose > 0:
         print >> sys.stderr, "Running on port %s"%options.port
diff --git a/test/test_dump.py b/test/test_dump.py
index 9aee5a31..0d0a6219 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -19,13 +19,22 @@ class uDumpMaster(libpry.AutoTree):
         assert o.verbosity == 2
         libpry.raises(AttributeError, dump.Options, nonexistent = 2)
 
-    def test_basic_verbosities(self):
+    def test_filter(self):
+            cs = StringIO()
+            o = dump.Options(
+                verbosity = 1
+            )
+            m = dump.DumpMaster(None, o, "~u foo", outfile=cs)
+            self._dummy_cycle(m)
+            assert not "GET" in cs.getvalue()
+
+    def test_basic(self):
         for i in (1, 2, 3):
             cs = StringIO()
             o = dump.Options(
                 verbosity = i
             )
-            m = dump.DumpMaster(None, o, cs)
+            m = dump.DumpMaster(None, o, "~s", outfile=cs)
             self._dummy_cycle(m)
             assert "GET" in cs.getvalue()
 
@@ -37,7 +46,7 @@ class uDumpMaster(libpry.AutoTree):
             verbosity = 0
         )
         cs = StringIO()
-        m = dump.DumpMaster(None, o, cs)
+        m = dump.DumpMaster(None, o, None, outfile=cs)
         self._dummy_cycle(m)
         del m
         assert len(list(flow.FlowReader(open(p)).stream())) == 1
@@ -48,7 +57,7 @@ class uDumpMaster(libpry.AutoTree):
             verbosity = 0
         )
         cs = StringIO()
-        libpry.raises(dump.DumpError, dump.DumpMaster, None, o, cs)
+        libpry.raises(dump.DumpError, dump.DumpMaster, None, o, None)
 
 
 
-- 
cgit v1.2.3