Add file writing to mitmdump.

This commit is contained in:
Aldo Cortesi 2011-02-17 10:18:38 +13:00
parent 7758385ac1
commit f0f1fb4b55
6 changed files with 54 additions and 8 deletions

View File

@ -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):

View File

@ -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):

View File

@ -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)

View File

@ -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)

View File

@ -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 = [

View File

@ -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()