streamfile: make options less crazy

Instead of a (file, mode) tuple, move to "streamfile" as the path, and
"streamfile_append" to specify if we should append or not.

Leave the command-line options intact for now.
This commit is contained in:
Aldo Cortesi 2016-11-03 15:12:04 +13:00
parent 9f77c79227
commit 77e6dfe35c
8 changed files with 19 additions and 56 deletions

View File

@ -33,10 +33,11 @@ class StreamFile:
if self.stream: if self.stream:
self.done() self.done()
if options.streamfile: if options.streamfile:
path, mode = options.streamfile if options.streamfile_append:
if mode not in ("wb", "ab"): mode = "ab"
raise exceptions.OptionsError("Invalid mode.") else:
self.start_stream_to_path(path, mode, self.filt) mode = "wb"
self.start_stream_to_path(options.streamfile, mode, self.filt)
def tcp_start(self, flow): def tcp_start(self, flow):
if self.stream: if self.stream:

View File

@ -48,7 +48,8 @@ class Options(optmanager.OptManager):
stream_large_bodies: Optional[int] = None, stream_large_bodies: Optional[int] = None,
verbosity: int = 2, verbosity: int = 2,
default_contentview: str = "auto", default_contentview: str = "auto",
streamfile: Optional[Tuple[str, str]] = None, streamfile: Optional[str] = None,
streamfile_append: bool = False,
server_replay_ignore_content: bool = False, server_replay_ignore_content: bool = False,
server_replay_ignore_params: Sequence[str] = (), server_replay_ignore_params: Sequence[str] = (),
server_replay_ignore_payload_params: Sequence[str] = (), server_replay_ignore_payload_params: Sequence[str] = (),
@ -109,6 +110,7 @@ class Options(optmanager.OptManager):
self.verbosity = verbosity self.verbosity = verbosity
self.default_contentview = default_contentview self.default_contentview = default_contentview
self.streamfile = streamfile self.streamfile = streamfile
self.streamfile_append = streamfile_append
self.server_replay_ignore_content = server_replay_ignore_content self.server_replay_ignore_content = server_replay_ignore_content
self.server_replay_ignore_params = server_replay_ignore_params self.server_replay_ignore_params = server_replay_ignore_params
self.server_replay_ignore_payload_params = server_replay_ignore_payload_params self.server_replay_ignore_payload_params = server_replay_ignore_payload_params

View File

@ -228,7 +228,8 @@ def get_common_options(args):
stickyauth=stickyauth, stickyauth=stickyauth,
stream_large_bodies=stream_large_bodies, stream_large_bodies=stream_large_bodies,
showhost=args.showhost, showhost=args.showhost,
streamfile=args.streamfile, streamfile=args.streamfile[0] if args.streamfile else None,
streamfile_append=True if args.streamfile and args.streamfile[1] == "a" else False,
verbosity=args.verbose, verbosity=args.verbose,
server_replay_nopop=args.server_replay_nopop, server_replay_nopop=args.server_replay_nopop,
server_replay_ignore_content=args.server_replay_ignore_content, server_replay_ignore_content=args.server_replay_ignore_content,
@ -342,12 +343,12 @@ def basic_options(parser):
streamfile = parser.add_mutually_exclusive_group() streamfile = parser.add_mutually_exclusive_group()
streamfile.add_argument( streamfile.add_argument(
"-w", "--wfile", "-w", "--wfile",
action="store", dest="streamfile", type=lambda f: (f, "wb"), action="store", dest="streamfile", type=lambda f: (f, "w"),
help="Write flows to file." help="Write flows to file."
) )
streamfile.add_argument( streamfile.add_argument(
"-a", "--afile", "-a", "--afile",
action="store", dest="streamfile", type=lambda f: (f, "ab"), action="store", dest="streamfile", type=lambda f: (f, "a"),
help="Append flows to file." help="Append flows to file."
) )
parser.add_argument( parser.add_argument(

View File

@ -399,7 +399,7 @@ class FlowListBox(urwid.ListBox):
signals.status_prompt_path.send( signals.status_prompt_path.send(
self, self,
prompt="Stream flows to", prompt="Stream flows to",
callback= lambda path: self.master.options.update(streamfile=(path, "ab")) callback= lambda path: self.master.options.update(streamfile=path)
) )
else: else:
return urwid.ListBox.keypress(self, size, key) return urwid.ListBox.keypress(self, size, key)

View File

@ -224,7 +224,7 @@ class StatusBar(urwid.WidgetWrap):
r.append("cripts:%s]" % len(self.master.options.scripts)) r.append("cripts:%s]" % len(self.master.options.scripts))
if self.master.options.streamfile: if self.master.options.streamfile:
r.append("[W:%s]" % self.master.options.streamfile[0]) r.append("[W:%s]" % self.master.options.streamfile)
return r return r

View File

@ -111,15 +111,6 @@ class WebMaster(master.Master):
"error" "error"
) )
if options.streamfile:
err = self.start_stream_to_path(
options.streamfile[0],
options.streamfile[1]
)
if err:
print("Stream file error: {}".format(err), file=sys.stderr)
sys.exit(1)
def _sig_add(self, view, flow): def _sig_add(self, view, flow):
app.ClientConnection.broadcast( app.ClientConnection.broadcast(
type="UPDATE_FLOWS", type="UPDATE_FLOWS",

View File

@ -16,15 +16,11 @@ def test_configure():
p = os.path.join(tdir, "foo") p = os.path.join(tdir, "foo")
tutils.raises( tutils.raises(
exceptions.OptionsError, exceptions.OptionsError,
tctx.configure, sa, streamfile=(tdir, "ab") tctx.configure, sa, streamfile=tdir
) )
tutils.raises( tutils.raises(
"invalid filter", "invalid filter",
tctx.configure, sa, streamfile=(p, "ab"), filtstr="~~" tctx.configure, sa, streamfile=p, filtstr="~~"
)
tutils.raises(
"invalid mode",
tctx.configure, sa, streamfile=(p, "xx")
) )
@ -38,7 +34,7 @@ def test_tcp():
with taddons.context() as tctx: with taddons.context() as tctx:
with tutils.tmpdir() as tdir: with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo") p = os.path.join(tdir, "foo")
tctx.configure(sa, streamfile=(p, "wb")) tctx.configure(sa, streamfile=p)
tt = tflow.ttcpflow() tt = tflow.ttcpflow()
sa.tcp_start(tt) sa.tcp_start(tt)
@ -53,7 +49,7 @@ def test_simple():
with tutils.tmpdir() as tdir: with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo") p = os.path.join(tdir, "foo")
tctx.configure(sa, streamfile=(p, "wb")) tctx.configure(sa, streamfile=p)
f = tflow.tflow(resp=True) f = tflow.tflow(resp=True)
sa.request(f) sa.request(f)
@ -61,7 +57,7 @@ def test_simple():
tctx.configure(sa, streamfile=None) tctx.configure(sa, streamfile=None)
assert rd(p)[0].response assert rd(p)[0].response
tctx.configure(sa, streamfile=(p, "ab")) tctx.configure(sa, streamfile=p, streamfile_append=True)
f = tflow.tflow() f = tflow.tflow()
sa.request(f) sa.request(f)
tctx.configure(sa, streamfile=None) tctx.configure(sa, streamfile=None)

View File

@ -2,7 +2,6 @@ from mitmproxy.test import tflow
import os import os
import io import io
import mitmproxy.io
from mitmproxy.tools import dump from mitmproxy.tools import dump
from mitmproxy import exceptions from mitmproxy import exceptions
from mitmproxy import proxy from mitmproxy import proxy
@ -126,33 +125,6 @@ class TestDumpMaster(mastertest.MasterTest):
f = self.cycle(m, b"content") f = self.cycle(m, b"content")
assert f.request.headers["one"] == "two" assert f.request.headers["one"] == "two"
def test_write(self):
with tutils.tmpdir() as d:
p = os.path.join(d, "a")
self.dummy_cycle(
self.mkmaster(None, outfile=(p, "wb"), verbosity=0), 1, b""
)
assert len(list(mitmproxy.io.FlowReader(open(p, "rb")).stream())) == 1
def test_write_append(self):
with tutils.tmpdir() as d:
p = os.path.join(d, "a.append")
self.dummy_cycle(
self.mkmaster(None, outfile=(p, "wb"), verbosity=0),
1, b""
)
self.dummy_cycle(
self.mkmaster(None, outfile=(p, "ab"), verbosity=0),
1, b""
)
assert len(list(mitmproxy.io.FlowReader(open(p, "rb")).stream())) == 2
def test_write_err(self):
tutils.raises(
exceptions.OptionsError,
self.mkmaster, None, outfile = ("nonexistentdir/foo", "wb")
)
def test_script(self): def test_script(self):
ret = self.dummy_cycle( ret = self.dummy_cycle(
self.mkmaster( self.mkmaster(