From 9b1f40da370860afbf75a34ba437092413e872ec Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Mon, 6 Mar 2017 13:42:11 +1300
Subject: [PATCH] Options unification: streamfile

We now have one option to control this. If the path is prefixed with a "+" we
append, otherwise we overwrite.
---
 mitmproxy/addons/streamfile.py           |  6 ++++--
 mitmproxy/options.py                     |  6 ++++--
 mitmproxy/tools/cmdline.py               | 27 ++----------------------
 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 5517e9dc3..624297f2b 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 59c57f3ee..c0ac3d679 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 2290086cf..d48fc737c 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 89dc2af3c..4105c1fc9 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)