codekingpro/portable-devtools

1"""2This module manages and invokes typed commands.3"""4 5import functools6import inspect7import logging8import sys9import textwrap10import types11from collections.abc import Callable12from collections.abc import Iterable13from collections.abc import Sequence14from typing import Any15from typing import NamedTuple16 17import pyparsing18 19import mitmproxy.types20from mitmproxy import command_lexer21from mitmproxy import exceptions22from mitmproxy.command_lexer import unquote23 24 25def verify_arg_signature(f: Callable, args: Iterable[Any], kwargs: dict) -> None:26    sig = inspect.signature(f, eval_str=True)27    try:28        sig.bind(*args, **kwargs)29    except TypeError as v:30        raise exceptions.CommandError("command argument mismatch: %s" % v.args[0])31 32 33def typename(t: type) -> str:34    """35    Translates a type to an explanatory string.36    """37    if t == inspect._empty:  # type: ignore38        raise exceptions.CommandError("missing type annotation")39    to = mitmproxy.types.CommandTypes.get(t, None)40    if not to:41        raise exceptions.CommandError(42            "unsupported type: %s" % getattr(t, "__name__", t)43        )44    return to.display45 46 47def _empty_as_none(x: Any) -> Any:48    if x == inspect.Signature.empty:49        return None50    return x51 52 53class CommandParameter(NamedTuple):54    name: str55    type: type56    kind: inspect._ParameterKind = inspect.Parameter.POSITIONAL_OR_KEYWORD57 58    def __str__(self):59        if self.kind is inspect.Parameter.VAR_POSITIONAL:60            return f"*{self.name}"61        else:62            return self.name63 64 65class Command:66    name: str67    manager: "CommandManager"68    signature: inspect.Signature69    help: str | None70 71    def __init__(self, manager: "CommandManager", name: str, func: Callable) -> None:72        self.name = name73        self.manager = manager74        self.func = func75        self.signature = inspect.signature(self.func, eval_str=True)76 77        if func.__doc__:78            txt = func.__doc__.strip()79            self.help = "\n".join(textwrap.wrap(txt))80        else:81            self.help = None82 83        # This fails with a CommandException if types are invalid84        for name, parameter in self.signature.parameters.items():85            t = parameter.annotation86            if not mitmproxy.types.CommandTypes.get(parameter.annotation, None):87                raise exceptions.CommandError(88                    f"Argument {name} has an unknown type {t} in {func}."89                )90        if self.return_type and not mitmproxy.types.CommandTypes.get(91            self.return_type, None92        ):93            raise exceptions.CommandError(94                f"Return type has an unknown type ({self.return_type}) in {func}."95            )96 97    @property98    def return_type(self) -> type | None:99        return _empty_as_none(self.signature.return_annotation)100 101    @property102    def parameters(self) -> list[CommandParameter]:103        """Returns a list of CommandParameters."""104        ret = []105        for name, param in self.signature.parameters.items():106            ret.append(CommandParameter(name, param.annotation, param.kind))107        return ret108 109    def signature_help(self) -> str:110        params = " ".join(str(param) for param in self.parameters)111        if self.return_type:112            ret = f" -> {typename(self.return_type)}"113        else:114            ret = ""115        return f"{self.name} {params}{ret}"116 117    def prepare_args(self, args: Sequence[str]) -> inspect.BoundArguments:118        try:119            bound_arguments = self.signature.bind(*args)120        except TypeError:121            expected = f"Expected: {self.signature.parameters}"122            received = f"Received: {args}"123            raise exceptions.CommandError(124                f"Command argument mismatch: \n    {expected}\n    {received}"125            )126 127        for name, value in bound_arguments.arguments.items():128            param = self.signature.parameters[name]129            convert_to = param.annotation130            if param.kind == param.VAR_POSITIONAL:131                bound_arguments.arguments[name] = tuple(132                    parsearg(self.manager, x, convert_to) for x in value133                )134            else:135                bound_arguments.arguments[name] = parsearg(136                    self.manager, value, convert_to137                )138 139        bound_arguments.apply_defaults()140 141        return bound_arguments142 143    def call(self, args: Sequence[str]) -> Any:144        """145        Call the command with a list of arguments. At this point, all146        arguments are strings.147        """148        bound_args = self.prepare_args(args)149        ret = self.func(*bound_args.args, **bound_args.kwargs)150        if ret is None and self.return_type is None:151            return152        typ = mitmproxy.types.CommandTypes.get(self.return_type)153        assert typ154        if not typ.is_valid(self.manager, typ, ret):155            raise exceptions.CommandError(156                f"{self.name} returned unexpected data - expected {typ.display}"157            )158        return ret159 160 161class ParseResult(NamedTuple):162    value: str163    type: type164    valid: bool165 166 167class CommandManager:168    commands: dict[str, Command]169 170    def __init__(self, master):171        self.master = master172        self.commands = {}173 174    def collect_commands(self, addon):175        for i in dir(addon):176            if not i.startswith("__"):177                o = getattr(addon, i)178                try:179                    # hasattr is not enough, see https://github.com/mitmproxy/mitmproxy/issues/3794180                    is_command = isinstance(getattr(o, "command_name", None), str)181                except Exception:182                    pass  # getattr may raise if o implements __getattr__.183                else:184                    if is_command:185                        try:186                            self.add(o.command_name, o)187                        except exceptions.CommandError as e:188                            logging.warning(189                                f"Could not load command {o.command_name}: {e}"190                            )191 192    def add(self, path: str, func: Callable):193        self.commands[path] = Command(self, path, func)194 195    @functools.lru_cache(maxsize=128)196    def parse_partial(197        self, cmdstr: str198    ) -> tuple[Sequence[ParseResult], Sequence[CommandParameter]]:199        """200        Parse a possibly partial command. Return a sequence of ParseResults and a sequence of remainder type help items.201        """202 203        parts: pyparsing.ParseResults = command_lexer.expr.parseString(204            cmdstr, parseAll=True205        )206 207        parsed: list[ParseResult] = []208        next_params: list[CommandParameter] = [209            CommandParameter("", mitmproxy.types.Cmd),210            CommandParameter("", mitmproxy.types.CmdArgs),211        ]212        expected: CommandParameter | None = None213        for part in parts:214            if part.isspace():215                parsed.append(216                    ParseResult(217                        value=part,218                        type=mitmproxy.types.Space,219                        valid=True,220                    )221                )222                continue223 224            if expected and expected.kind is inspect.Parameter.VAR_POSITIONAL:225                assert not next_params226            elif next_params:227                expected = next_params.pop(0)228            else:229                expected = CommandParameter("", mitmproxy.types.Unknown)230 231            arg_is_known_command = (232                expected.type == mitmproxy.types.Cmd and part in self.commands233            )234            arg_is_unknown_command = (235                expected.type == mitmproxy.types.Cmd and part not in self.commands236            )237            command_args_following = (238                next_params and next_params[0].type == mitmproxy.types.CmdArgs239            )240            if arg_is_known_command and command_args_following:241                next_params = self.commands[part].parameters + next_params[1:]242            if arg_is_unknown_command and command_args_following:243                next_params.pop(0)244 245            to = mitmproxy.types.CommandTypes.get(expected.type, None)246            valid = False247            if to:248                try:249                    to.parse(self, expected.type, part)250                except ValueError:251                    valid = False252                else:253                    valid = True254 255            parsed.append(256                ParseResult(257                    value=part,258                    type=expected.type,259                    valid=valid,260                )261            )262 263        return parsed, next_params264 265    def call(self, command_name: str, *args: Any) -> Any:266        """267        Call a command with native arguments. May raise CommandError.268        """269        if command_name not in self.commands:270            raise exceptions.CommandError("Unknown command: %s" % command_name)271        return self.commands[command_name].func(*args)272 273    def call_strings(self, command_name: str, args: Sequence[str]) -> Any:274        """275        Call a command using a list of string arguments. May raise CommandError.276        """277        if command_name not in self.commands:278            raise exceptions.CommandError("Unknown command: %s" % command_name)279 280        return self.commands[command_name].call(args)281 282    def execute(self, cmdstr: str) -> Any:283        """284        Execute a command string. May raise CommandError.285        """286        parts, _ = self.parse_partial(cmdstr)287        if not parts:288            raise exceptions.CommandError(f"Invalid command: {cmdstr!r}")289        command_name, *args = (290            unquote(part.value) for part in parts if part.type != mitmproxy.types.Space291        )292        return self.call_strings(command_name, args)293 294    def dump(self, out=sys.stdout) -> None:295        cmds = list(self.commands.values())296        cmds.sort(key=lambda x: x.signature_help())297        for c in cmds:298            for hl in (c.help or "").splitlines():299                print("# " + hl, file=out)300            print(c.signature_help(), file=out)301            print(file=out)302 303 304def parsearg(manager: CommandManager, spec: str, argtype: type) -> Any:305    """306    Convert a string to a argument to the appropriate type.307    """308    t = mitmproxy.types.CommandTypes.get(argtype, None)309    if not t:310        raise exceptions.CommandError(f"Unsupported argument type: {argtype}")311    try:312        return t.parse(manager, argtype, spec)313    except ValueError as e:314        raise exceptions.CommandError(str(e)) from e315 316 317def command(name: str | None = None):318    def decorator(function):319        @functools.wraps(function)320        def wrapper(*args, **kwargs):321            verify_arg_signature(function, args, kwargs)322            return function(*args, **kwargs)323 324        wrapper.__dict__["command_name"] = name or function.__name__.replace("_", ".")325        return wrapper326 327    return decorator328 329 330def argument(name, type):331    """332    Set the type of a command argument at runtime. This is useful for more333    specific types such as mitmproxy.types.Choice, which we cannot annotate334    directly as mypy does not like that.335    """336 337    def decorator(f: types.FunctionType) -> types.FunctionType:338        assert name in f.__annotations__339        f.__annotations__[name] = type340        return f341 342    return decorator343