From f5511350eb383a2d2d73c6d1ae328f6eb00f501f Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Wed, 16 Feb 2011 23:03:46 +1300
Subject: [PATCH] Tweaks and cleanups for mitmdump.

---
 libmproxy/dump.py | 64 +++++++++++++++++++++++++++--------------------
 mitmdump          |  7 +++++-
 test/test_dump.py | 11 ++++++--
 3 files changed, 52 insertions(+), 30 deletions(-)

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 039fb9a74..76bf3ae2a 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -1,10 +1,21 @@
 import sys
 import flow
 
+class Options(object):
+    __slots__ = [
+        "verbosity",
+        "wfile",
+    ]
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+
 class DumpMaster(flow.FlowMaster):
-    def __init__(self, server, verbosity, outfile=sys.stderr):
-        self.verbosity, self.outfile = verbosity, outfile
+    def __init__(self, server, options, outfile=sys.stdout):
         flow.FlowMaster.__init__(self, server, flow.State())
+        self.outfile = outfile
+        self.o = options
 
     def handle_clientconnection(self, r):
         flow.FlowMaster.handle_clientconnection(self, r)
@@ -20,31 +31,30 @@ class DumpMaster(flow.FlowMaster):
 
     def handle_response(self, msg):
         f = flow.FlowMaster.handle_response(self, msg)
-        if not f:
-            return
-        if 0 < self.verbosity < 3:
-            print >> self.outfile, ">>",
-            print >> self.outfile, msg.request.short()
-        if self.verbosity == 1:
-            print >> self.outfile, "<<",
-            print >> self.outfile, msg.short()
-        elif self.verbosity == 2:
-            print >> self.outfile, "<<"
-            for i in msg.assemble().splitlines():
-                print >> self.outfile, "\t", i
-            print >> self.outfile, "<<"
-        elif self.verbosity == 3:
-            print >> self.outfile, ">>"
-            for i in msg.request.assemble().splitlines():
-                print >> self.outfile, "\t", i
-            print >> self.outfile, ">>"
-            print >> self.outfile, "<<"
-            for i in msg.assemble().splitlines():
-                print >> self.outfile, "\t", i
-            print >> self.outfile, "<<"
-        msg.ack()
-        self.state.delete_flow(f)
-        print [i.get_state() for i in self.state.flow_list]
+        if f:
+            if 0 < self.o.verbosity < 3:
+                print >> self.outfile, ">>",
+                print >> self.outfile, msg.request.short()
+            if self.o.verbosity == 1:
+                print >> self.outfile, "<<",
+                print >> self.outfile, msg.short()
+            elif self.o.verbosity == 2:
+                print >> self.outfile, "<<"
+                for i in msg.assemble().splitlines():
+                    print >> self.outfile, "\t", i
+                print >> self.outfile, "<<"
+            elif self.o.verbosity == 3:
+                print >> self.outfile, ">>"
+                for i in msg.request.assemble().splitlines():
+                    print >> self.outfile, "\t", i
+                print >> self.outfile, ">>"
+                print >> self.outfile, "<<"
+                for i in msg.assemble().splitlines():
+                    print >> self.outfile, "\t", i
+                print >> self.outfile, "<<"
+
+            msg.ack()
+            self.state.delete_flow(f)
 
 
 # begin nocover
diff --git a/mitmdump b/mitmdump
index 11b715689..c0ed7dfe8 100755
--- a/mitmdump
+++ b/mitmdump
@@ -76,7 +76,12 @@ if __name__ == '__main__':
         ciphers = options.ciphers
     )
     server = proxy.ProxyServer(options.port)
-    m = dump.DumpMaster(server, options.verbose)
+
+    dumpopts = dump.Options(
+        verbosity = options.verbose
+    )
+    m = dump.DumpMaster(server, dumpopts)
+
     if options.verbose > 0:
         print >> sys.stderr, "Running on port %s"%options.port
     m.run()
diff --git a/test/test_dump.py b/test/test_dump.py
index da9de2b47..f9a6bdf40 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -4,7 +4,6 @@ from libmproxy import dump
 import utils
 
 
-
 class uDumpMaster(libpry.AutoTree):
     def _dummy_cycle(self, m):
         req = utils.treq()
@@ -14,10 +13,18 @@ class uDumpMaster(libpry.AutoTree):
         m.handle_request(req)
         m.handle_response(resp)
 
+    def test_options(self):
+        o = dump.Options(verbosity = 2)
+        assert o.verbosity == 2
+        libpry.raises(AttributeError, dump.Options, nonexistent = 2)
+
     def test_basic_verbosities(self):
         for i in (1, 2, 3):
             cs = StringIO()
-            m = dump.DumpMaster(None, i, cs)
+            o = dump.Options(
+                verbosity = i
+            )
+            m = dump.DumpMaster(None, o, cs)
             self._dummy_cycle(m)
             assert "GET" in cs.getvalue()