Initial port of mitmdump to Flows.

3 files changed, 74 insertions, 31 deletions

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 1fe1c095..7eff4992 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -1,40 +1,52 @@
 import sys
-import controller
+import flow
 
-#begin nocover
-class DumpMaster(controller.Master):
-    """
-        A simple master that just dumps to screen.
-    """
-    def __init__(self, server, verbosity):
-        self.verbosity = verbosity
-        controller.Master.__init__(self, server)
+class DumpMaster(flow.FlowMaster):
+    def __init__(self, server, verbosity, outfile=sys.stderr):
+        self.verbosity, self.outfile = verbosity, outfile
+        flow.FlowMaster.__init__(self, server, flow.State())
 
-    def run(self):
-        try:
-            return controller.Master.run(self)
-        except KeyboardInterrupt:
-            self.shutdown()
+    def handle_clientconnection(self, r):
+        flow.FlowMaster.handle_clientconnection(self, r)
+        r.ack()
+
+    def handle_error(self, r):
+        flow.FlowMaster.handle_error(self, r)
+        r.ack()
+
+    def handle_request(self, r):
+        flow.FlowMaster.handle_request(self, r)
+        r.ack()
 
     def handle_response(self, msg):
+        f = flow.FlowMaster.handle_response(self, msg)
         if 0 < self.verbosity < 3:
-            print >> sys.stderr, ">>",
-            print >> sys.stderr, msg.request.short()
+            print >> self.outfile, ">>",
+            print >> self.outfile, msg.request.short()
         if self.verbosity == 1:
-            print >> sys.stderr, "<<",
-            print >> sys.stderr, msg.short()
+            print >> self.outfile, "<<",
+            print >> self.outfile, msg.short()
         elif self.verbosity == 2:
-            print >> sys.stderr, "<<"
+            print >> self.outfile, "<<"
             for i in msg.assemble().splitlines():
-                print >> sys.stderr, "\t", i
-            print >> sys.stderr, "<<"
+                print >> self.outfile, "\t", i
+            print >> self.outfile, "<<"
         elif self.verbosity == 3:
-            print >> sys.stderr, ">>"
+            print >> self.outfile, ">>"
             for i in msg.request.assemble().splitlines():
-                print >> sys.stderr, "\t", i
-            print >> sys.stderr, ">>"
-            print >> sys.stderr, "<<"
+                print >> self.outfile, "\t", i
+            print >> self.outfile, ">>"
+            print >> self.outfile, "<<"
             for i in msg.assemble().splitlines():
-                print >> sys.stderr, "\t", i
-            print >> sys.stderr, "<<"
+                print >> self.outfile, "\t", i
+            print >> self.outfile, "<<"
         msg.ack()
+
+
+# begin nocover
+    def run(self):
+        try:
+            return flow.FlowMaster.run(self)
+        except KeyboardInterrupt:
+            self.shutdown()
+
diff --git a/test/test_console.py b/test/test_console.py
index 179b116d..93312824 100644
--- a/test/test_console.py
+++ b/test/test_console.py
@@ -105,11 +105,12 @@ class uformat_keyvals(libpry.AutoTree):
 class uformat_flow(libpry.AutoTree):
     def test_simple(self):
         f = utils.tflow()
-        assert ('focus', '>> ') not in console.format_flow(f, False)
-        assert ('focus', '>> ') in console.format_flow(f, True)
+        foc = ('focus', '>>')
+        assert foc not in console.format_flow(f, False)
+        assert foc in console.format_flow(f, True)
 
-        assert ('focus', '>> ') not in console.format_flow(f, False, True)
-        assert ('focus', '>> ') in console.format_flow(f, True, True)
+        assert foc not in console.format_flow(f, False, True)
+        assert foc in console.format_flow(f, True, True)
 
         f.response = utils.tresp()
         f.request = f.response.request
diff --git a/test/test_dump.py b/test/test_dump.py
new file mode 100644
index 00000000..da9de2b4
--- /dev/null
+++ b/test/test_dump.py
@@ -0,0 +1,30 @@
+from cStringIO import StringIO
+import libpry
+from libmproxy import dump
+import utils
+
+
+
+class uDumpMaster(libpry.AutoTree):
+    def _dummy_cycle(self, m):
+        req = utils.treq()
+        cc = req.client_conn
+        resp = utils.tresp(req)
+        m.handle_clientconnection(cc)
+        m.handle_request(req)
+        m.handle_response(resp)
+
+    def test_basic_verbosities(self):
+        for i in (1, 2, 3):
+            cs = StringIO()
+            m = dump.DumpMaster(None, i, cs)
+            self._dummy_cycle(m)
+            assert "GET" in cs.getvalue()
+
+
+
+tests = [
+    uDumpMaster()
+]
+
+