codekingpro/portable-devtools

1import logging2import os3from collections import defaultdict4from collections.abc import Sequence5from functools import cache6 7import ruamel.yaml.error8 9import mitmproxy.types10from mitmproxy import command11from mitmproxy import ctx12from mitmproxy import exceptions13from mitmproxy.tools.console import commandexecutor14from mitmproxy.tools.console import signals15 16 17class KeyBindingError(Exception):18    pass19 20 21Contexts = {22    "chooser",23    "commands",24    "commonkey",25    "dataviewer",26    "eventlog",27    "flowlist",28    "flowview",29    "global",30    "grideditor",31    "help",32    "keybindings",33    "options",34}35 36 37navkeys = [38    "m_start",39    "m_end",40    "m_next",41    "m_select",42    "up",43    "down",44    "page_up",45    "page_down",46    "left",47    "right",48]49 50 51class Binding:52    def __init__(self, key, command, contexts, help):53        self.key, self.command, self.contexts = key, command, sorted(contexts)54        self.help = help55 56    def keyspec(self):57        """58        Translate the key spec from a convenient user specification to one59        Urwid understands.60        """61        return self.key.replace("space", " ")62 63    def key_short(self) -> str:64        return (65            self.key.replace("enter", "⏎").replace("right", "→").replace("space", "␣")66        )67 68    def sortkey(self):69        return self.key + ",".join(self.contexts)70 71 72class Keymap:73    def __init__(self, master):74        self.executor = commandexecutor.CommandExecutor(master)75        self.keys: dict[str, dict[str, Binding]] = defaultdict(dict)76        self.bindings = []77 78    def _check_contexts(self, contexts):79        if not contexts:80            raise ValueError("Must specify at least one context.")81        for c in contexts:82            if c not in Contexts:83                raise ValueError("Unsupported context: %s" % c)84 85    def _on_change(self) -> None:86        signals.keybindings_change.send()87        self.binding_for_help.cache_clear()88 89    def add(self, key: str, command: str, contexts: Sequence[str], help="") -> None:90        """91        Add a key to the key map.92        """93        self._check_contexts(contexts)94 95        for b in self.bindings:96            if b.key == key and b.command.strip() == command.strip():97                b.contexts = sorted(list(set(b.contexts + contexts)))98                if help:99                    b.help = help100                self.bind(b)101                break102        else:103            self.remove(key, contexts)104            b = Binding(key=key, command=command, contexts=contexts, help=help)105            self.bindings.append(b)106            self.bind(b)107        self._on_change()108 109    def remove(self, key: str, contexts: Sequence[str]) -> None:110        """111        Remove a key from the key map.112        """113        self._check_contexts(contexts)114        for c in contexts:115            b = self.get(c, key)116            if b:117                self.unbind(b)118                b.contexts = [x for x in b.contexts if x != c]119                if b.contexts:120                    self.bindings.append(b)121                    self.bind(b)122        self._on_change()123 124    def bind(self, binding: Binding) -> None:125        for c in binding.contexts:126            self.keys[c][binding.keyspec()] = binding127 128    def unbind(self, binding: Binding) -> None:129        """130        Unbind also removes the binding from the list.131        """132        for c in binding.contexts:133            del self.keys[c][binding.keyspec()]134            self.bindings = [b for b in self.bindings if b != binding]135        self._on_change()136 137    def get(self, context: str, key: str) -> Binding | None:138        if context in self.keys:139            return self.keys[context].get(key, None)140        return None141 142    @cache143    def binding_for_help(self, help: str) -> Binding | None:144        for b in self.bindings:145            if b.help == help:146                return b147        return None148 149    def list(self, context: str) -> Sequence[Binding]:150        b = [x for x in self.bindings if context in x.contexts or context == "all"]151        single = [x for x in b if len(x.key.split()) == 1]152        multi = [x for x in b if len(x.key.split()) != 1]153        single.sort(key=lambda x: x.sortkey())154        multi.sort(key=lambda x: x.sortkey())155        return single + multi156 157    def handle(self, context: str, key: str) -> str | None:158        """159        Returns the key if it has not been handled, or None.160        """161        b = self.get(context, key) or self.get("global", key)162        if b:163            self.executor(b.command)164            return None165        return key166 167    def handle_only(self, context: str, key: str) -> str | None:168        """169        Like handle, but ignores global bindings. Returns the key if it has170        not been handled, or None.171        """172        b = self.get(context, key)173        if b:174            self.executor(b.command)175            return None176        return key177 178 179keyAttrs = {180    "key": lambda x: isinstance(x, str),181    "cmd": lambda x: isinstance(x, str),182    "ctx": lambda x: isinstance(x, list) and [isinstance(v, str) for v in x],183    "help": lambda x: isinstance(x, str),184}185requiredKeyAttrs = {"key", "cmd"}186 187 188class KeymapConfig:189    defaultFile = "keys.yaml"190 191    def __init__(self, master):192        self.master = master193 194    @command.command("console.keymap.load")195    def keymap_load_path(self, path: mitmproxy.types.Path) -> None:196        try:197            self.load_path(self.master.keymap, path)  # type: ignore198        except (OSError, KeyBindingError) as e:199            raise exceptions.CommandError("Could not load key bindings - %s" % e) from e200 201    def running(self):202        p = os.path.join(os.path.expanduser(ctx.options.confdir), self.defaultFile)203        if os.path.exists(p):204            try:205                self.load_path(self.master.keymap, p)206            except KeyBindingError as e:207                logging.error(e)208 209    def load_path(self, km, p):210        if os.path.exists(p) and os.path.isfile(p):211            with open(p, encoding="utf8") as f:212                try:213                    txt = f.read()214                except UnicodeDecodeError as e:215                    raise KeyBindingError(f"Encoding error - expected UTF8: {p}: {e}")216            try:217                vals = self.parse(txt)218            except KeyBindingError as e:219                raise KeyBindingError(f"Error reading {p}: {e}") from e220            for v in vals:221                user_ctxs = v.get("ctx", ["global"])222                try:223                    km._check_contexts(user_ctxs)224                    km.remove(v["key"], user_ctxs)225                    km.add(226                        key=v["key"],227                        command=v["cmd"],228                        contexts=user_ctxs,229                        help=v.get("help", None),230                    )231                except ValueError as e:232                    raise KeyBindingError(f"Error reading {p}: {e}") from e233 234    def parse(self, text):235        try:236            data = ruamel.yaml.YAML(typ="safe", pure=True).load(text)237        except ruamel.yaml.error.MarkedYAMLError as v:238            if hasattr(v, "problem_mark"):239                snip = v.problem_mark.get_snippet()240                raise KeyBindingError(241                    "Key binding config error at line %s:\n%s\n%s"242                    % (v.problem_mark.line + 1, snip, v.problem)243                )244            else:245                raise KeyBindingError("Could not parse key bindings.")246        if not data:247            return []248        if not isinstance(data, list):249            raise KeyBindingError("Invalid keybinding config - expected a list of keys")250 251        for k in data:252            unknown = k.keys() - keyAttrs.keys()253            if unknown:254                raise KeyBindingError("Unknown key attributes: %s" % unknown)255            missing = requiredKeyAttrs - k.keys()256            if missing:257                raise KeyBindingError("Missing required key attributes: %s" % unknown)258            for attr in k.keys():259                if not keyAttrs[attr](k[attr]):260                    raise KeyBindingError("Invalid type for %s" % attr)261 262        return data263