diff --git a/libmproxy/console.py b/libmproxy/console.py index ac1ef16cc..c14f1aefb 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 76bf3ae2a..e30c24229 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 7172b4db4..c9be74836 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 2d379f378..47fe7c2d1 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 f9a6bdf40..9aee5a313 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 7354061b3..9871b090e 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()