codekingpro/portable-devtools

1import asyncio2import logging3 4from . import ctx as mitmproxy_ctx5from .addons import termlog6from .proxy.mode_specs import ReverseMode7from .utils import asyncio_utils8from mitmproxy import addonmanager9from mitmproxy import command10from mitmproxy import eventsequence11from mitmproxy import hooks12from mitmproxy import http13from mitmproxy import log14from mitmproxy import options15 16logger = logging.getLogger(__name__)17 18 19class Master:20    """21    The master handles mitmproxy's main event loop.22    """23 24    event_loop: asyncio.AbstractEventLoop25    _termlog_addon: termlog.TermLog | None = None26 27    def __init__(28        self,29        opts: options.Options | None,30        event_loop: asyncio.AbstractEventLoop | None = None,31        with_termlog: bool = False,32    ):33        self.options: options.Options = opts or options.Options()34        self.commands = command.CommandManager(self)35        self.addons = addonmanager.AddonManager(self)36 37        if with_termlog:38            self._termlog_addon = termlog.TermLog()39            self.addons.add(self._termlog_addon)40 41        self.log = log.Log(self)  # deprecated, do not use.42        self._legacy_log_events = log.LegacyLogEvents(self)43        self._legacy_log_events.install()44 45        # We expect an active event loop here already because some addons46        # may want to spawn tasks during the initial configuration phase,47        # which happens before run().48        self.event_loop = event_loop or asyncio.get_running_loop()49        self.should_exit = asyncio.Event()50        mitmproxy_ctx.master = self51        mitmproxy_ctx.log = self.log  # deprecated, do not use.52        mitmproxy_ctx.options = self.options53 54    async def run(self) -> None:55        with (56            asyncio_utils.install_exception_handler(self._asyncio_exception_handler),57            asyncio_utils.set_eager_task_factory(),58        ):59            self.should_exit.clear()60 61            # Can we exit before even bringing up servers?62            if ec := self.addons.get("errorcheck"):63                await ec.shutdown_if_errored()64            if ps := self.addons.get("proxyserver"):65                # This may block for some proxy modes, so we also monitor should_exit.66                await asyncio.wait(67                    [68                        asyncio_utils.create_task(69                            ps.setup_servers(), name="setup_servers", keep_ref=False70                        ),71                        asyncio_utils.create_task(72                            self.should_exit.wait(), name="should_exit", keep_ref=False73                        ),74                    ],75                    return_when=asyncio.FIRST_COMPLETED,76                )77                if self.should_exit.is_set():78                    return79                # Did bringing up servers fail?80                if ec := self.addons.get("errorcheck"):81                    await ec.shutdown_if_errored()82 83            try:84                await self.running()85                # Any errors in the final part of startup?86                if ec := self.addons.get("errorcheck"):87                    await ec.shutdown_if_errored()88                    ec.finish()89 90                await self.should_exit.wait()91            finally:92                # if running() was called, we also always want to call done().93                # .wait might be cancelled (e.g. by sys.exit), so  this needs to be in a finally block.94                await self.done()95 96    def shutdown(self):97        """98        Shut down the proxy. This method is thread-safe.99        """100        # We may add an exception argument here.101        self.event_loop.call_soon_threadsafe(self.should_exit.set)102 103    async def running(self) -> None:104        await self.addons.trigger_event(hooks.RunningHook())105 106    async def done(self) -> None:107        await self.addons.trigger_event(hooks.DoneHook())108        self._legacy_log_events.uninstall()109        if self._termlog_addon is not None:110            self._termlog_addon.uninstall()111 112    def _asyncio_exception_handler(self, loop, context) -> None:113        try:114            exc: Exception = context["exception"]115        except KeyError:116            logger.error(f"Unhandled asyncio error: {context}")117        else:118            if isinstance(exc, OSError) and exc.errno == 10038:119                return  # suppress https://bugs.python.org/issue43253120            logger.error(121                "Unhandled error in task.",122                exc_info=(type(exc), exc, exc.__traceback__),123            )124 125    async def load_flow(self, f):126        """127        Loads a flow128        """129 130        if (131            isinstance(f, http.HTTPFlow)132            and len(self.options.mode) == 1133            and self.options.mode[0].startswith("reverse:")134        ):135            # When we load flows in reverse proxy mode, we adjust the target host to136            # the reverse proxy destination for all flows we load. This makes it very137            # easy to replay saved flows against a different host.138            # We may change this in the future so that clientplayback always replays to the first mode.139            mode = ReverseMode.parse(self.options.mode[0])140            assert isinstance(mode, ReverseMode)141            f.request.host, f.request.port, *_ = mode.address142            f.request.scheme = mode.scheme143 144        for e in eventsequence.iterate(f):145            await self.addons.handle_lifecycle(e)146