Add file writing to mitmdump.

6 files changed, 54 insertions, 8 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index ac1ef16c..c14f1aef 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -822,7 +822,7 @@ class ConsoleMaster(flow.FlowMaster):
         self.stickycookie = None
         self.stickyhosts = {}
 
-        if options.cache is not None:
+        if getattr(options, "cache", None) is not None:
             self.state.start_recording(recorder.Recorder(options))
 
     def spawn_external_viewer(self, data, contenttype):
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 76bf3ae2..e30c2422 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -1,6 +1,9 @@
-import sys
+import sys, os
 import flow
 
+class DumpError(Exception): pass
+
+
 class Options(object):
     __slots__ = [
         "verbosity",
@@ -9,6 +12,9 @@ class Options(object):
     def __init__(self, **kwargs):
         for k, v in kwargs.items():
             setattr(self, k, v)
+        for i in self.__slots__:
+            if not hasattr(self, i):
+                setattr(self, i, None)
 
 
 class DumpMaster(flow.FlowMaster):
@@ -17,6 +23,14 @@ class DumpMaster(flow.FlowMaster):
         self.outfile = outfile
         self.o = options
 
+        if options.wfile:
+            path = os.path.expanduser(options.wfile)
+            try:
+                f = file(path, "wb")
+                self.fwriter = flow.FlowWriter(f)
+            except IOError, v:
+                raise DumpError(v.strerror)
+
     def handle_clientconnection(self, r):
         flow.FlowMaster.handle_clientconnection(self, r)
         r.ack()
@@ -56,6 +70,9 @@ class DumpMaster(flow.FlowMaster):
             msg.ack()
             self.state.delete_flow(f)
 
+            if self.o.wfile:
+                self.fwriter.add(f)
+
 
 # begin nocover
     def run(self):
diff --git a/libmproxy/utils.py b/libmproxy/utils.py
index 7172b4db..c9be7483 100644
--- a/libmproxy/utils.py
+++ b/libmproxy/utils.py
@@ -525,6 +525,6 @@ def process_certificate_option_group(options):
         options.certpath = os.path.expanduser(options.certpath)
     elif options.cacert is not None:
         options.certpath = os.path.dirname(options.cacert)
-    if options.cache is not None:
+    if getattr(options, "cache", None) is not None:
         options.cache = os.path.expanduser(options.cache)
 
diff --git a/mitmdump b/mitmdump
index 2d379f37..47fe7c2d 100755
--- a/mitmdump
+++ b/mitmdump
@@ -38,6 +38,9 @@ if __name__ == '__main__':
     parser.add_option("-v", "--verbose",
                       action="count", dest="verbose", default=1,
                       help="Increase verbosity. Can be passed multiple times.")
+    parser.add_option("-w", "--writefile",
+                      action="store", dest="wfile", default=None,
+                      help="Write flows to file.")
     options, args = parser.parse_args()
 
     if options.quiet:
@@ -54,7 +57,8 @@ if __name__ == '__main__':
     server = proxy.ProxyServer(options.port)
 
     dumpopts = dump.Options(
-        verbosity = options.verbose
+        verbosity = options.verbose,
+        wfile = options.wfile
     )
     m = dump.DumpMaster(server, dumpopts)
 
diff --git a/test/test_dump.py b/test/test_dump.py
index f9a6bdf4..9aee5a31 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -1,6 +1,7 @@
+import os
 from cStringIO import StringIO
 import libpry
-from libmproxy import dump
+from libmproxy import dump, flow
 import utils
 
 
@@ -28,6 +29,30 @@ class uDumpMaster(libpry.AutoTree):
             self._dummy_cycle(m)
             assert "GET" in cs.getvalue()
 
+    def test_write(self):
+        d = self.tmpdir()
+        p = os.path.join(d, "a")
+        o = dump.Options(
+            wfile = p,
+            verbosity = 0
+        )
+        cs = StringIO()
+        m = dump.DumpMaster(None, o, cs)
+        self._dummy_cycle(m)
+        del m
+        assert len(list(flow.FlowReader(open(p)).stream())) == 1
+
+    def test_write_err(self):
+        o = dump.Options(
+            wfile = "nonexistentdir/foo",
+            verbosity = 0
+        )
+        cs = StringIO()
+        libpry.raises(dump.DumpError, dump.DumpMaster, None, o, cs)
+
+
+
+
 
 
 tests = [
diff --git a/test/test_flow.py b/test/test_flow.py
index 7354061b..9871b090 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -141,12 +141,12 @@ class uState(libpry.AutoTree):
         assert c.lookup(req)
 
         newreq = utils.treq()
-        assert not c.add_request(newreq)
-        assert not c.lookup(newreq)
+        assert c.add_request(newreq)
+        assert c.lookup(newreq)
 
         resp = utils.tresp(req)
         assert c.add_response(resp)
-        assert len(c.flow_list) == 1
+        assert len(c.flow_list) == 2
         assert c.lookup(resp)
 
         newresp = utils.tresp()