codekingpro/portable-devtools

1from __future__ import annotations2 3import argparse4import asyncio5import logging6import os7import signal8import sys9from collections.abc import Callable10from collections.abc import Sequence11from typing import Any12from typing import TypeVar13 14from mitmproxy import exceptions15from mitmproxy import master16from mitmproxy import options17from mitmproxy import optmanager18from mitmproxy.tools import cmdline19from mitmproxy.utils import arg_check20from mitmproxy.utils import debug21 22 23def process_options(parser, opts, args):24    if args.version:25        print(debug.dump_system_info())26        sys.exit(0)27    if args.quiet or args.options or args.commands:28        # also reduce log verbosity if --options or --commands is passed,29        # we don't want log messages from regular startup then.30        args.termlog_verbosity = "error"31        args.flow_detail = 032    if args.verbose:33        args.termlog_verbosity = "debug"34        args.flow_detail = 235 36    adict = {37        key: val for key, val in vars(args).items() if key in opts and val is not None38    }39    opts.update(**adict)40 41 42T = TypeVar("T", bound=master.Master)43 44 45def run(46    master_cls: type[T],47    make_parser: Callable[[options.Options], argparse.ArgumentParser],48    arguments: Sequence[str],49    extra: Callable[[Any], dict] | None = None,50) -> T:  # pragma: no cover51    """52    extra: Extra argument processing callable which returns a dict of53    options.54    """55 56    async def main() -> T:57        logging.getLogger().setLevel(logging.DEBUG)58        logging.getLogger("tornado").setLevel(logging.WARNING)59        logging.getLogger("asyncio").setLevel(logging.WARNING)60        logging.getLogger("hpack").setLevel(logging.WARNING)61        logging.getLogger("urwid").setLevel(logging.INFO)62        logging.getLogger("quic").setLevel(63            logging.WARNING64        )  # aioquic uses a different prefix...65        debug.register_info_dumpers()66 67        opts = options.Options()68        master = master_cls(opts)69 70        parser = make_parser(opts)71 72        # To make migration from 2.x to 3.0 bearable.73        if "-R" in sys.argv and sys.argv[sys.argv.index("-R") + 1].startswith("http"):74            print(75                "To use mitmproxy in reverse mode please use --mode reverse:SPEC instead"76            )77 78        try:79            args = parser.parse_args(arguments)80        except SystemExit:81            arg_check.check()82            sys.exit(1)83 84        try:85            opts.set(*args.setoptions, defer=True)86            optmanager.load_paths(87                opts,88                os.path.join(opts.confdir, "config.yaml"),89                os.path.join(opts.confdir, "config.yml"),90            )91            process_options(parser, opts, args)92 93            if args.options:94                # Load custom addons so that their options are registered95                if sl := master.addons.get("scriptloader"):96                    for s in sl.addons:97                        if s.ns is None:98                            s.loadscript()99                optmanager.dump_defaults(opts, sys.stdout)100                sys.exit(0)101            if args.commands:102                master.commands.dump()103                sys.exit(0)104            if extra:105                if args.filter_args:106                    logging.info(107                        f'Only processing flows that match "{" & ".join(args.filter_args)}"'108                    )109                opts.update(**extra(args))110 111        except exceptions.OptionsError as e:112            print(f"{sys.argv[0]}: {e}", file=sys.stderr)113            sys.exit(1)114 115        loop = asyncio.get_running_loop()116 117        def _sigint(*_):118            loop.call_soon_threadsafe(119                getattr(master, "prompt_for_exit", master.shutdown)120            )121 122        def _sigterm(*_):123            loop.call_soon_threadsafe(master.shutdown)124 125        try:126            # Prefer loop.add_signal_handler where it is available127            # https://github.com/mitmproxy/mitmproxy/issues/7128128            loop.add_signal_handler(signal.SIGINT, _sigint)129            loop.add_signal_handler(signal.SIGTERM, _sigterm)130        except NotImplementedError:131            # Fall back to `signal.signal` for platforms where that is not available (Windows' Proactorloop)132            signal.signal(signal.SIGINT, _sigint)133            signal.signal(signal.SIGTERM, _sigterm)134 135        # to fix the issue mentioned https://github.com/mitmproxy/mitmproxy/issues/6744136        # by setting SIGPIPE to SIG_IGN, the process will not terminate and continue to run137        if hasattr(signal, "SIGPIPE"):138            signal.signal(signal.SIGPIPE, signal.SIG_IGN)139 140        await master.run()141        return master142 143    return asyncio.run(main())144 145 146def mitmproxy(args=None) -> int | None:  # pragma: no cover147    from mitmproxy.tools import console148 149    run(console.master.ConsoleMaster, cmdline.mitmproxy, args)150    return None151 152 153def mitmdump(args=None) -> int | None:  # pragma: no cover154    from mitmproxy.tools import dump155 156    def extra(args):157        if args.filter_args:158            v = " ".join(args.filter_args)159            return dict(160                save_stream_filter=v,161                readfile_filter=v,162                dumper_filter=v,163            )164        return {}165 166    run(dump.DumpMaster, cmdline.mitmdump, args, extra)167    return None168 169 170def mitmweb(args=None) -> int | None:  # pragma: no cover171    from mitmproxy.tools import web172 173    run(web.master.WebMaster, cmdline.mitmweb, args)174    return None175