Flatten commandline arg structure, extract common run func for tools

2 files changed, 51 insertions, 100 deletions

diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py
index 13577997..e789aeab 100644
--- a/mitmproxy/tools/cmdline.py
+++ b/mitmproxy/tools/cmdline.py
@@ -8,7 +8,7 @@ from mitmproxy import version
 CONFIG_PATH = os.path.join(options.CA_DIR, "config.yaml")
 
 
-def basic_options(parser, opts):
+def common_options(parser, opts):
     parser.add_argument(
         '--version',
         action='store_true',
@@ -20,30 +20,38 @@ def basic_options(parser, opts):
         help="show program's short version number and exit",
         version=version.VERSION
     )
-    opts.make_parser(parser, "anticache")
-    opts.make_parser(parser, "cadir")
-    opts.make_parser(parser, "showhost")
     parser.add_argument(
         "-q", "--quiet",
         action="store_true", dest="quiet",
         help="Quiet."
     )
-    opts.make_parser(parser, "rfile")
-    opts.make_parser(parser, "scripts", metavar="SCRIPT")
-    opts.make_parser(parser, "stickycookie", metavar="FILTER")
-    opts.make_parser(parser, "stickyauth", metavar="FILTER")
     parser.add_argument(
         "-v", "--verbose",
         action="store_const", dest="verbose", const=3,
         help="Increase log verbosity."
     )
+    parser.add_argument(
+        "--conf",
+        type=str, dest="conf", default=CONFIG_PATH,
+        metavar="PATH",
+        help="Configuration file"
+    )
+
+    # Basic options
+    opts.make_parser(parser, "mode")
+    opts.make_parser(parser, "anticache")
+    opts.make_parser(parser, "cadir")
+    opts.make_parser(parser, "showhost")
+    opts.make_parser(parser, "rfile")
+    opts.make_parser(parser, "scripts", metavar="SCRIPT")
+    opts.make_parser(parser, "stickycookie", metavar="FILTER")
+    opts.make_parser(parser, "stickyauth", metavar="FILTER")
     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")
 
-
-def proxy_options(parser, opts):
+    # Proxy options
     group = parser.add_argument_group("Proxy Options")
     opts.make_parser(group, "listen_host")
     opts.make_parser(group, "ignore_hosts", metavar="HOST")
@@ -64,9 +72,7 @@ def proxy_options(parser, opts):
     opts.make_parser(group, "upstream_bind_address", metavar="ADDR")
     opts.make_parser(group, "keep_host_header")
 
-
-def proxy_ssl_options(parser, opts):
-    # TODO: Agree to consistently either use "upstream" or "server".
+    # Proxy SSL options
     group = parser.add_argument_group("SSL")
     opts.make_parser(group, "certs", metavar="SPEC")
     opts.make_parser(group, "ciphers_server", metavar="CIPHERS")
@@ -80,20 +86,17 @@ def proxy_ssl_options(parser, opts):
     opts.make_parser(group, "ssl_version_client", metavar="VERSION")
     opts.make_parser(group, "ssl_version_server", metavar="VERSION")
 
-
-def onboarding_app(parser, opts):
+    # Onboarding app
     group = parser.add_argument_group("Onboarding App")
     opts.make_parser(group, "onboarding")
     opts.make_parser(group, "onboarding_host", metavar="HOST")
     opts.make_parser(group, "onboarding_port", metavar="PORT")
 
-
-def client_replay(parser, opts):
+    # Client replay
     group = parser.add_argument_group("Client Replay")
     opts.make_parser(group, "client_replay", metavar="PATH")
 
-
-def server_replay(parser, opts):
+    # Server replay
     group = parser.add_argument_group("Server Replay")
     opts.make_parser(group, "server_replay", metavar="PATH")
     opts.make_parser(group, "replay_kill_extra")
@@ -107,8 +110,7 @@ def server_replay(parser, opts):
     opts.make_parser(payload, "server_replay_ignore_params")
     opts.make_parser(payload, "server_replay_ignore_host")
 
-
-def replacements(parser, opts):
+    # Replacements
     group = parser.add_argument_group(
         "Replacements",
         """
@@ -120,8 +122,7 @@ def replacements(parser, opts):
     opts.make_parser(group, "replacements", metavar="PATTERN")
     opts.make_parser(group, "replacement_files", metavar="PATTERN")
 
-
-def set_headers(parser, opts):
+    # Set headers
     group = parser.add_argument_group(
         "Set Headers",
         """
@@ -132,8 +133,7 @@ def set_headers(parser, opts):
     )
     opts.make_parser(group, "setheaders", metavar="PATTERN")
 
-
-def proxy_authentication(parser, opts):
+    # Proxy authentication
     group = parser.add_argument_group(
         "Proxy Authentication",
         """
@@ -146,25 +146,6 @@ def proxy_authentication(parser, opts):
     opts.make_parser(group, "auth_htpasswd", metavar="PATH")
 
 
-def common_options(parser, opts):
-    parser.add_argument(
-        "--conf",
-        type=str, dest="conf", default=CONFIG_PATH,
-        metavar="PATH",
-        help="Configuration file"
-    )
-    opts.make_parser(parser, "mode")
-    basic_options(parser, opts)
-    proxy_options(parser, opts)
-    proxy_ssl_options(parser, opts)
-    onboarding_app(parser, opts)
-    client_replay(parser, opts)
-    server_replay(parser, opts)
-    replacements(parser, opts)
-    set_headers(parser, opts)
-    proxy_authentication(parser, opts)
-
-
 def mitmproxy(opts):
     parser = argparse.ArgumentParser(usage="%(prog)s [options]")
     common_options(parser, opts)
diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py
index 05877e2c..0478d3c4 100644
--- a/mitmproxy/tools/main.py
+++ b/mitmproxy/tools/main.py
@@ -47,7 +47,6 @@ def process_options(parser, options, args):
             adict[n] = getattr(args, n)
     options.merge(adict)
 
-    debug.register_info_dumpers()
     pconf = config.ProxyConfig(options)
     if options.no_server:
         return server.DummyServer(pconf)
@@ -59,46 +58,18 @@ def process_options(parser, options, args):
             sys.exit(1)
 
 
-def mitmproxy(args=None):  # pragma: no cover
-    if os.name == "nt":
-        print("Error: mitmproxy's console interface is not supported on Windows. "
-              "You can run mitmdump or mitmweb instead.", file=sys.stderr)
-        sys.exit(1)
-    from mitmproxy.tools import console
-
-    version_check.check_pyopenssl_version()
-    assert_utf8_env()
-
-    console_options = options.Options()
-    parser = cmdline.mitmproxy(console_options)
-    args = parser.parse_args(args)
-
-    try:
-        console_options.load_paths(args.conf)
-        server = process_options(parser, console_options, args)
-        m = console.master.ConsoleMaster(console_options, server)
-    except exceptions.OptionsError as e:
-        print("mitmproxy: %s" % e, file=sys.stderr)
-        sys.exit(1)
-    try:
-        m.run()
-    except (KeyboardInterrupt, RuntimeError):
-        pass
-
-
-def mitmdump(args=None):  # pragma: no cover
-    from mitmproxy.tools import dump
-
+def run(MasterKlass, args):  # pragma: no cover
     version_check.check_pyopenssl_version()
+    debug.register_info_dumpers()
 
-    dump_options = options.Options()
-    parser = cmdline.mitmdump(dump_options)
+    opts = options.Options()
+    parser = cmdline.mitmdump(opts)
     args = parser.parse_args(args)
     master = None
     try:
-        dump_options.load_paths(args.conf)
-        server = process_options(parser, dump_options, args)
-        master = dump.DumpMaster(dump_options, server)
+        opts.load_paths(args.conf)
+        server = process_options(parser, opts, args)
+        master = MasterKlass(opts, server)
 
         def cleankill(*args, **kwargs):
             master.shutdown()
@@ -106,32 +77,31 @@ def mitmdump(args=None):  # pragma: no cover
         signal.signal(signal.SIGTERM, cleankill)
         master.run()
     except exceptions.OptionsError as e:
-        print("mitmdump: %s" % e, file=sys.stderr)
+        print("%s: %s" % (sys.argv[0], e), file=sys.stderr)
         sys.exit(1)
     except (KeyboardInterrupt, RuntimeError):
         pass
-    if master is None or master.has_errored:
-        print("mitmdump: errors occurred during run", file=sys.stderr)
+    if master is None or getattr(master, "has_errored", None):
+        print("%s: errors occurred during run" % sys.argv[0], file=sys.stderr)
         sys.exit(1)
 
 
-def mitmweb(args=None):  # pragma: no cover
-    from mitmproxy.tools import web
+def mitmproxy(args=None):  # pragma: no cover
+    if os.name == "nt":
+        print("Error: mitmproxy's console interface is not supported on Windows. "
+              "You can run mitmdump or mitmweb instead.", file=sys.stderr)
+        sys.exit(1)
+    assert_utf8_env()
 
-    version_check.check_pyopenssl_version()
+    from mitmproxy.tools import console
+    run(console.master.ConsoleMaster, args)
 
-    web_options = options.Options()
-    parser = cmdline.mitmweb(web_options)
-    args = parser.parse_args(args)
 
-    try:
-        web_options.load_paths(args.conf)
-        server = process_options(parser, web_options, args)
-        m = web.master.WebMaster(web_options, server)
-    except exceptions.OptionsError as e:
-        print("mitmweb: %s" % e, file=sys.stderr)
-        sys.exit(1)
-    try:
-        m.run()
-    except (KeyboardInterrupt, RuntimeError):
-        pass
+def mitmdump(args=None):  # pragma: no cover
+    from mitmproxy.tools import dump
+    run(dump.DumpMaster, args)
+
+
+def mitmweb(args=None):  # pragma: no cover
+    from mitmproxy.tools import web
+    run(web.master.WebMaster, args)
\ No newline at end of file