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:
self.done()
if options.streamfile:
path, mode = options.streamfile
if mode not in ("wb", "ab"):
raise exceptions.OptionsError("Invalid mode.")
self.start_stream_to_path(path, mode, self.filt)
if options.streamfile_append:
mode = "ab"
else:
mode = "wb"
self.start_stream_to_path(options.streamfile, mode, self.filt)
def tcp_start(self, flow):
if self.stream:

View File

@ -48,7 +48,8 @@ class Options(optmanager.OptManager):
stream_large_bodies: Optional[int] = None,
verbosity: int = 2,
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_params: Sequence[str] = (),
server_replay_ignore_payload_params: Sequence[str] = (),
@ -109,6 +110,7 @@ class Options(optmanager.OptManager):
self.verbosity = verbosity
self.default_contentview = default_contentview
self.streamfile = streamfile
self.streamfile_append = streamfile_append
self.server_replay_ignore_content = server_replay_ignore_content
self.server_replay_ignore_params = server_replay_ignore_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,
stream_large_bodies=stream_large_bodies,
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,
server_replay_nopop=args.server_replay_nopop,
server_replay_ignore_content=args.server_replay_ignore_content,
@ -342,12 +343,12 @@ def basic_options(parser):
streamfile = parser.add_mutually_exclusive_group()
streamfile.add_argument(
"-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."
)
streamfile.add_argument(
"-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."
)
parser.add_argument(

View File

@ -399,7 +399,7 @@ class FlowListBox(urwid.ListBox):
signals.status_prompt_path.send(
self,
prompt="Stream flows to",
callback= lambda path: self.master.options.update(streamfile=(path, "ab"))
callback= lambda path: self.master.options.update(streamfile=path)
)
else:
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))
if self.master.options.streamfile:
r.append("[W:%s]" % self.master.options.streamfile[0])
r.append("[W:%s]" % self.master.options.streamfile)
return r

View File

@ -111,15 +111,6 @@ class WebMaster(master.Master):
"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):
app.ClientConnection.broadcast(
type="UPDATE_FLOWS",

View File

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

View File

@ -2,7 +2,6 @@ from mitmproxy.test import tflow
import os
import io
import mitmproxy.io
from mitmproxy.tools import dump
from mitmproxy import exceptions
from mitmproxy import proxy
@ -126,33 +125,6 @@ class TestDumpMaster(mastertest.MasterTest):
f = self.cycle(m, b"content")
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):
ret = self.dummy_cycle(
self.mkmaster(