pm3py/pm3py/cli/rawcli/completer.py

"""ipython-style completion for rawcli.

Completes the leading token in two ways, both driven by the **identified transponder**:

- by name — control verbs and the transponder's catalog commands (help text as the tooltip);
- by **opcode** — as you type hex or binary bytes, the transponder's command opcodes are matched
  against the input (respecting a ``0x``/``0b`` prefix and the current entry mode), so ``60`` and
  ``01100000`` both surface ``GET_VERSION``.

``help <cmd>`` completes catalog command names too. Paired in the app with menu completion while
typing and history auto-suggestion for the ipython feel.
"""
from __future__ import annotations

import re

from prompt_toolkit.completion import Completer, Completion

from .parser import CONTROL_VERBS
from .catalog import catalog_for
from .memory import landmark_pages

_HEXDIGITS = set("0123456789abcdefABCDEF")
_BINDIGITS = set("01")

# NAME( <first-arg-being-typed> — matches only while the first argument is under the cursor
# (a comma ends the match, so second args aren't touched).
_CALL_ARG = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*)\(\s*([^,)]*)$")


def _opcode(tc) -> int | None:
    """The command's first byte (opcode), obtained by building it with placeholder args. Returns
    None for commands with no ``build`` (read-only credential ops that carry no opcode)."""
    if tc.build is None:
        return None
    try:
        built = tc.build(*(["0"] * len(tc.params)))
        return built[0] if built else None
    except Exception:
        return None


class RawCompleter(Completer):
    def __init__(self, session):
        self._session = session

    def get_completions(self, document, complete_event):
        text = document.text_before_cursor
        stripped = text.lstrip()

        # inside READ(<page> / WRITE(<block> — offer the tag's memory map for the page argument
        m = _CALL_ARG.match(stripped)
        if m:
            yield from self._page_arg(m.group(1), m.group(2))
            return

        tokens = stripped.split()
        at_word = bool(tokens) and not stripped.endswith(" ")
        word = tokens[-1] if at_word else ""

        completing_first = (" " not in stripped) or (len(tokens) == 1 and at_word)
        if completing_first:
            yield from self._first_token(word)
        elif tokens and tokens[0].lower() == "help":
            yield from self._catalog_names(word)

    def _page_arg(self, cmd_name: str, partial: str):
        """Complete a page/block argument with the identified tag's memory landmarks, rendered in
        the active entry mode — ``0x04`` in hex, ``0b00000100`` in binary (every bit visible,
        since for config pages each bit carries its own meaning). ``READ(`` lists the whole map;
        typing narrows it. A ``0x``/``0b`` prefix on the argument overrides the session mode."""
        catalog = catalog_for(self._session)
        if catalog is None:
            return
        tc = catalog.get(cmd_name)
        if tc is None or not tc.params or tc.params[0] not in ("page", "block"):
            return
        mode = self._arg_mode(partial)
        body = partial.lower()
        if body.startswith(("0x", "0b")):
            body = body[2:]
        for page, role in landmark_pages(self._session.transponder):
            if mode == "bin":
                val, forms = f"0b{page:08b}", (f"{page:08b}", f"{page:b}")
            else:
                val, forms = f"0x{page:02X}", (f"{page:02x}", f"{page:x}", str(page))
            if body and not any(f.startswith(body) for f in forms):
                continue
            yield Completion(val, start_position=-len(partial),
                             display=f"{val}  {role}", display_meta=role)

    def _arg_mode(self, partial: str) -> str:
        """hex/bin for an argument token — a 0x/0b prefix wins, else the session entry mode."""
        low = partial.lower()
        if low.startswith("0x"):
            return "hex"
        if low.startswith("0b"):
            return "bin"
        return "bin" if getattr(self._session, "entry_mode", "hex") == "bin" else "hex"

    def _first_token(self, word: str):
        low = word.lower()
        for verb in sorted(CONTROL_VERBS):
            if verb.startswith(low):
                yield Completion(verb, start_position=-len(word), display_meta="command")

        catalog = catalog_for(self._session)
        if catalog is None:
            return

        # by command name
        for name in catalog.names():
            if name.lower().startswith(low):
                tc = catalog.get(name)
                yield Completion(f"{name}(", start_position=-len(word),
                                 display=tc.usage(), display_meta=tc.help)

        # by opcode, from hex/binary input (transponder-dependent)
        yield from self._opcode_matches(word, catalog)

    def _numeric(self, word: str) -> tuple[str, str]:
        """(digits, base) for an entry token — honours a 0x/0b prefix, else the session mode."""
        low = word.lower()
        if low.startswith("0x"):
            return low[2:], "hex"
        if low.startswith("0b"):
            return low[2:], "bin"
        return low, getattr(self._session, "entry_mode", "hex")

    def _opcode_matches(self, word: str, catalog):
        digits, base = self._numeric(word)
        if not digits:
            return
        if base == "hex":
            valid, render = _HEXDIGITS, (lambda op: f"{op:02x}")
        else:
            valid, render = _BINDIGITS, (lambda op: f"{op:08b}")
        if any(ch not in valid for ch in digits):
            return
        low = word.lower()
        for name in catalog.names():
            if name.lower().startswith(low):
                continue                              # already offered by name
            tc = catalog.get(name)
            op = _opcode(tc)
            if op is not None and render(op).startswith(digits):
                yield Completion(f"{name}(", start_position=-len(word),
                                 display=f"{name} (0x{op:02X})",
                                 display_meta=f"opcode 0x{op:02X} — {tc.help}")

    def _catalog_names(self, word: str):
        catalog = catalog_for(self._session)
        if catalog is None:
            return
        low = word.lower()
        for name in catalog.names():
            if name.lower().startswith(low):
                yield Completion(name, start_position=-len(word),
                                 display_meta=catalog.get(name).help)