From 3d9c2233be328ecd1c243ddb6a180a0f1a3ef6c5 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Tue, 7 Mar 2017 10:51:24 +1300 Subject: [PATCH] Flatten commandline arg structure, extract common run func for tools --- mitmproxy/tools/cmdline.py | 69 ++++++++++++------------------- mitmproxy/tools/main.py | 84 ++++++++++++-------------------------- 2 files changed, 52 insertions(+), 101 deletions(-) diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index 13577997d..e789aeab4 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 05877e2c9..0478d3c44 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 - +def run(MasterKlass, args): # pragma: no cover version_check.check_pyopenssl_version() - assert_utf8_env() + debug.register_info_dumpers() - 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 - - version_check.check_pyopenssl_version() - - 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 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() + + from mitmproxy.tools import console + run(console.master.ConsoleMaster, args) + + +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 - - version_check.check_pyopenssl_version() - - 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 + run(web.master.WebMaster, args) \ No newline at end of file