diff options
| -rw-r--r-- | mitmproxy/addons/streamfile.py | 9 | ||||
| -rw-r--r-- | mitmproxy/options.py | 4 | ||||
| -rw-r--r-- | mitmproxy/tools/cmdline.py | 7 | ||||
| -rw-r--r-- | mitmproxy/tools/console/flowlist.py | 2 | ||||
| -rw-r--r-- | mitmproxy/tools/console/statusbar.py | 2 | ||||
| -rw-r--r-- | mitmproxy/tools/web/master.py | 9 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_streamfile.py | 14 | ||||
| -rw-r--r-- | test/mitmproxy/test_dump.py | 28 |
8 files changed, 19 insertions, 56 deletions
diff --git a/mitmproxy/addons/streamfile.py b/mitmproxy/addons/streamfile.py index 76befeb2..377f277d 100644 --- a/mitmproxy/addons/streamfile.py +++ b/mitmproxy/addons/streamfile.py @@ -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: diff --git a/mitmproxy/options.py b/mitmproxy/options.py index 84aa4a74..49791400 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -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 diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index cb26378d..debe6db9 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -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( diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py index fc8368e7..6ad7f656 100644 --- a/mitmproxy/tools/console/flowlist.py +++ b/mitmproxy/tools/console/flowlist.py @@ -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) diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py index ebf2bca8..e3424493 100644 --- a/mitmproxy/tools/console/statusbar.py +++ b/mitmproxy/tools/console/statusbar.py @@ -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 diff --git a/mitmproxy/tools/web/master.py b/mitmproxy/tools/web/master.py index 3e1b5933..2f1fd4e3 100644 --- a/mitmproxy/tools/web/master.py +++ b/mitmproxy/tools/web/master.py @@ -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", diff --git a/test/mitmproxy/addons/test_streamfile.py b/test/mitmproxy/addons/test_streamfile.py index f63f29c4..82a4345b 100644 --- a/test/mitmproxy/addons/test_streamfile.py +++ b/test/mitmproxy/addons/test_streamfile.py @@ -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) diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py index aa3228e4..e331637d 100644 --- a/test/mitmproxy/test_dump.py +++ b/test/mitmproxy/test_dump.py @@ -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( |