codekingpro/portable-devtools

1import contextlib2import inspect3import logging4import pprint5import sys6import traceback7import types8from collections.abc import Callable9from collections.abc import Sequence10from dataclasses import dataclass11from typing import Any12 13from mitmproxy import exceptions14from mitmproxy import flow15from mitmproxy import hooks16 17logger = logging.getLogger(__name__)18 19 20def _get_name(itm):21    return getattr(itm, "name", itm.__class__.__name__.lower())22 23 24def cut_traceback(tb, func_name):25    """26    Cut off a traceback at the function with the given name.27    The func_name's frame is excluded.28 29    Args:30        tb: traceback object, as returned by sys.exc_info()[2]31        func_name: function name32 33    Returns:34        Reduced traceback.35    """36    tb_orig = tb37    for _, _, fname, _ in traceback.extract_tb(tb):38        tb = tb.tb_next39        if fname == func_name:40            break41    return tb or tb_orig42 43 44@contextlib.contextmanager45def safecall():46    try:47        yield48    except (exceptions.AddonHalt, exceptions.OptionsError):49        raise50    except Exception:51        etype, value, tb = sys.exc_info()52        tb = cut_traceback(tb, "invoke_addon_sync")53        tb = cut_traceback(tb, "invoke_addon")54        assert etype55        assert value56        logger.error(57            f"Addon error: {value}",58            exc_info=(etype, value, tb),59        )60 61 62class Loader:63    """64    A loader object is passed to the load() event when addons start up.65    """66 67    def __init__(self, master):68        self.master = master69 70    def add_option(71        self,72        name: str,73        typespec: type,74        default: Any,75        help: str,76        choices: Sequence[str] | None = None,77    ) -> None:78        """79        Add an option to mitmproxy.80 81        Help should be a single paragraph with no linebreaks - it will be82        reflowed by tools. Information on the data type should be omitted -83        it will be generated and added by tools as needed.84        """85        assert not isinstance(choices, str)86        if name in self.master.options:87            existing = self.master.options._options[name]88            same_signature = (89                existing.name == name90                and existing.typespec == typespec91                and existing.default == default92                and existing.help == help93                and existing.choices == choices94            )95            if same_signature:96                return97            else:98                logger.warning("Over-riding existing option %s" % name)99        self.master.options.add_option(name, typespec, default, help, choices)100 101    def add_command(self, path: str, func: Callable) -> None:102        """Add a command to mitmproxy.103 104        Unless you are generating commands programatically,105        this API should be avoided. Decorate your function with `@mitmproxy.command.command` instead.106        """107        self.master.commands.add(path, func)108 109 110def traverse(chain):111    """112    Recursively traverse an addon chain.113    """114    for a in chain:115        yield a116        if hasattr(a, "addons"):117            yield from traverse(a.addons)118 119 120@dataclass121class LoadHook(hooks.Hook):122    """123    Called when an addon is first loaded. This event receives a Loader124    object, which contains methods for adding options and commands. This125    method is where the addon configures itself.126    """127 128    loader: Loader129 130 131class AddonManager:132    def __init__(self, master):133        self.lookup = {}134        self.chain = []135        self.master = master136        master.options.changed.connect(self._configure_all)137 138    def _configure_all(self, updated):139        self.trigger(hooks.ConfigureHook(updated))140 141    def clear(self):142        """143        Remove all addons.144        """145        for a in self.chain:146            self.invoke_addon_sync(a, hooks.DoneHook())147        self.lookup = {}148        self.chain = []149 150    def get(self, name):151        """152        Retrieve an addon by name. Addon names are equal to the .name153        attribute on the instance, or the lower case class name if that154        does not exist.155        """156        return self.lookup.get(name, None)157 158    def register(self, addon):159        """160        Register an addon, call its load event, and then register all its161        sub-addons. This should be used by addons that dynamically manage162        addons.163 164        If the calling addon is already running, it should follow with165        running and configure events. Must be called within a current166        context.167        """168        api_changes = {169            # mitmproxy 6 -> mitmproxy 7170            "clientconnect": f"The clientconnect event has been removed, use client_connected instead",171            "clientdisconnect": f"The clientdisconnect event has been removed, use client_disconnected instead",172            "serverconnect": "The serverconnect event has been removed, use server_connect and server_connected instead",173            "serverdisconnect": f"The serverdisconnect event has been removed, use server_disconnected instead",174            # mitmproxy 8 -> mitmproxy 9175            "add_log": "The add_log event has been deprecated, use Python's builtin logging module instead",176        }177        for a in traverse([addon]):178            for old, msg in api_changes.items():179                if hasattr(a, old):180                    logger.warning(181                        f"{msg}. For more details, see https://docs.mitmproxy.org/dev/addons-api-changelog/."182                    )183            name = _get_name(a)184            if name in self.lookup:185                raise exceptions.AddonManagerError(186                    "An addon called '%s' already exists." % name187                )188        loader = Loader(self.master)189        self.invoke_addon_sync(addon, LoadHook(loader))190        for a in traverse([addon]):191            name = _get_name(a)192            self.lookup[name] = a193        for a in traverse([addon]):194            self.master.commands.collect_commands(a)195        self.master.options.process_deferred()196        return addon197 198    def add(self, *addons):199        """200        Add addons to the end of the chain, and run their load event.201        If any addon has sub-addons, they are registered.202        """203        for i in addons:204            self.chain.append(self.register(i))205 206    def remove(self, addon):207        """208        Remove an addon and all its sub-addons.209 210        If the addon is not in the chain - that is, if it's managed by a211        parent addon - it's the parent's responsibility to remove it from212        its own addons attribute.213        """214        for a in traverse([addon]):215            n = _get_name(a)216            if n not in self.lookup:217                raise exceptions.AddonManagerError("No such addon: %s" % n)218            self.chain = [i for i in self.chain if i is not a]219            del self.lookup[_get_name(a)]220        self.invoke_addon_sync(addon, hooks.DoneHook())221 222    def __len__(self):223        return len(self.chain)224 225    def __str__(self):226        return pprint.pformat([str(i) for i in self.chain])227 228    def __contains__(self, item):229        name = _get_name(item)230        return name in self.lookup231 232    async def handle_lifecycle(self, event: hooks.Hook):233        """234        Handle a lifecycle event.235        """236        message = event.args()[0]237 238        await self.trigger_event(event)239 240        if isinstance(message, flow.Flow):241            await self.trigger_event(hooks.UpdateHook([message]))242 243    def _iter_hooks(self, addon, event: hooks.Hook):244        """245        Enumerate all hook callables belonging to the given addon246        """247        assert isinstance(event, hooks.Hook)248        for a in traverse([addon]):249            func = getattr(a, event.name, None)250            if func:251                if callable(func):252                    yield a, func253                elif isinstance(func, types.ModuleType):254                    # we gracefully exclude module imports with the same name as hooks.255                    # For example, a user may have "from mitmproxy import log" in an addon,256                    # which has the same name as the "log" hook. In this particular case,257                    # we end up in an error loop because we "log" this error.258                    pass259                else:260                    raise exceptions.AddonManagerError(261                        f"Addon handler {event.name} ({a}) not callable"262                    )263 264    async def invoke_addon(self, addon, event: hooks.Hook):265        """266        Asynchronously invoke an event on an addon and all its children.267        """268        for addon, func in self._iter_hooks(addon, event):269            res = func(*event.args())270            # Support both async and sync hook functions271            if res is not None and inspect.isawaitable(res):272                await res273 274    def invoke_addon_sync(self, addon, event: hooks.Hook):275        """276        Invoke an event on an addon and all its children.277        """278        for addon, func in self._iter_hooks(addon, event):279            if inspect.iscoroutinefunction(func):280                raise exceptions.AddonManagerError(281                    f"Async handler {event.name} ({addon}) cannot be called from sync context"282                )283            func(*event.args())284 285    async def trigger_event(self, event: hooks.Hook):286        """287        Asynchronously trigger an event across all addons.288        """289        for i in self.chain:290            try:291                with safecall():292                    await self.invoke_addon(i, event)293            except exceptions.AddonHalt:294                return295 296    def trigger(self, event: hooks.Hook):297        """298        Trigger an event across all addons.299 300        This API is discouraged and may be deprecated in the future.301        Use `trigger_event()` instead, which provides the same functionality but supports async hooks.302        """303        for i in self.chain:304            try:305                with safecall():306                    self.invoke_addon_sync(i, event)307            except exceptions.AddonHalt:308                return309