mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 08:11:00 +00:00
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:
parent
9f77c79227
commit
77e6dfe35c
@ -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:
|
||||
|
@ -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
|
||||
|
@ -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(
|
||||
|
@ -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)
|
||||
|
@ -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
|
||||
|
||||
|
@ -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",
|
||||
|
@ -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)
|
||||
|
@ -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(
|
||||
|
Loading…
Reference in New Issue
Block a user