diff options
| -rw-r--r-- | mitmproxy/addons/streamfile.py | 6 | ||||
| -rw-r--r-- | mitmproxy/options.py | 6 | ||||
| -rw-r--r-- | mitmproxy/tools/cmdline.py | 27 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_streamfile.py | 2 |
4 files changed, 11 insertions, 30 deletions
diff --git a/mitmproxy/addons/streamfile.py b/mitmproxy/addons/streamfile.py index 5517e9dc..624297f2 100644 --- a/mitmproxy/addons/streamfile.py +++ b/mitmproxy/addons/streamfile.py @@ -35,11 +35,13 @@ class StreamFile: if self.stream: self.done() if options.streamfile: - if options.streamfile_append: + if options.streamfile.startswith("+"): + path = options.streamfile[1:] mode = "ab" else: + path = options.streamfile mode = "wb" - self.start_stream_to_path(options.streamfile, mode, self.filt) + self.start_stream_to_path(path, mode, self.filt) def tcp_start(self, flow): if self.stream: diff --git a/mitmproxy/options.py b/mitmproxy/options.py index 59c57f3e..c0ac3d67 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -142,8 +142,10 @@ class Options(optmanager.OptManager): "Log verbosity." ) self.add_option("default_contentview", "auto", str) - self.add_option("streamfile", None, Optional[str]) - self.add_option("streamfile_append", False, bool) + self.add_option( + "streamfile", None, Optional[str], + help="Write flows to file. Prefix path with + to append." + ) self.add_option( "server_replay_ignore_content", False, bool, "Ignore request's content while searching for a saved flow to replay." diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index 2290086c..d48fc737 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -15,18 +15,6 @@ class ParseException(Exception): def get_common_options(args): - if args.streamfile and args.streamfile[0] == args.rfile: - if args.streamfile[1] == "wb": - raise exceptions.OptionsError( - "Cannot use '{}' for both reading and writing flows. " - "Are you looking for --afile?".format(args.rfile) - ) - else: - raise exceptions.OptionsError( - "Cannot use '{}' for both reading and appending flows. " - "That would trigger an infinite loop." - ) - # Proxy config certs = [] for i in args.certs or []: @@ -96,8 +84,7 @@ def get_common_options(args): stickyauth=args.stickyauth, stream_large_bodies=args.stream_large_bodies, showhost=args.showhost, - streamfile=args.streamfile[0] if args.streamfile else None, - streamfile_append=True if args.streamfile and args.streamfile[1] == "a" else False, + streamfile=args.streamfile, verbosity=args.verbose, server_replay_nopop=args.server_replay_nopop, server_replay_ignore_content=args.server_replay_ignore_content, @@ -168,17 +155,7 @@ def basic_options(parser, opts): action="store_const", dest="verbose", const=3, help="Increase log verbosity." ) - streamfile = parser.add_mutually_exclusive_group() - streamfile.add_argument( - "-w", "--wfile", - 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, "a"), - help="Append flows to file." - ) + opts.make_parser(parser, "streamfile") opts.make_parser(parser, "anticomp") opts.make_parser(parser, "body_size_limit", metavar="SIZE") opts.make_parser(parser, "stream_large_bodies") diff --git a/test/mitmproxy/addons/test_streamfile.py b/test/mitmproxy/addons/test_streamfile.py index 89dc2af3..4105c1fc 100644 --- a/test/mitmproxy/addons/test_streamfile.py +++ b/test/mitmproxy/addons/test_streamfile.py @@ -59,7 +59,7 @@ def test_simple(): tctx.configure(sa, streamfile=None) assert rd(p)[0].response - tctx.configure(sa, streamfile=p, streamfile_append=True) + tctx.configure(sa, streamfile="+" + p) f = tflow.tflow() sa.request(f) tctx.configure(sa, streamfile=None) |