pm3py/pm3py/cli/clientcli/tree.py

"""Command-tree model for ``pm3py client``.

Loads the stock Proxmark3 client's machine-readable help dump (``doc/commands.json``) and builds a
prefix trie over the full command paths (``"hf mf rdbl"``). The JSON keys are **leaf** paths only —
interior nodes (``"hf"``, ``"hf mf"``) are synthesised from the path segments (no leaf path is a
prefix of another, so the trie is unambiguous). Option strings are pre-rendered by the client
(``"-k, --key <hex> key, 6 hex bytes"``); :func:`parse_option` splits them into
``(flags, metavar, description)`` for option-flag completion.

Regenerate the JSON for another client build with::

    proxmark3 --fulltext | client/pyscripts/pm3_help2json.py - -
"""
from __future__ import annotations

import json
import os
import re
from dataclasses import dataclass, field
from pathlib import Path

#: env override for the commands.json path (highest priority in :func:`load_command_tree`).
_ENV_VAR = "PM3_COMMANDS_JSON"

#: a single option flag, e.g. ``-a``, ``-1``, ``--key`` (numeric short flags like ``-1`` exist).
_FLAG = re.compile(r"^--?[A-Za-z0-9][\w-]*$")


def find_repo_root(start=None):
    """Locate a Proxmark3 client checkout (stock Iceman/RRG or a fork) by its ``doc/commands.json``
    marker, walking up from ``$PM3_REPO``/``$PM3_ROOT``, ``start``, the current directory, and this
    file's location. Returns a :class:`~pathlib.Path` to the repo root, or ``None``. This is what
    lets a dropped-in single file find whatever checkout it lives beside."""
    seeds = []
    for env in ("PM3_REPO", "PM3_ROOT"):
        if os.environ.get(env):
            seeds.append(Path(os.environ[env]))
    if start:
        seeds.append(Path(start))
    seeds.append(Path.cwd())
    try:
        seeds.append(Path(__file__).resolve().parent)
    except NameError:                                  # __file__ absent (e.g. frozen/REPL)
        pass
    seen = set()
    for seed in seeds:
        base = seed.resolve() if seed.exists() else seed
        for cand in (base, *base.parents):
            if cand in seen:
                continue
            seen.add(cand)
            if (cand / "doc" / "commands.json").is_file():
                return cand                            # dropped inside the checkout
            sub = cand / "proxmark3"                   # a side-by-side checkout (dev layout)
            if (sub / "doc" / "commands.json").is_file():
                return sub
    return None


@dataclass
class CommandNode:
    """A node in the command trie. ``entry`` is the ``commands.json`` leaf dict for a command, or
    ``None`` for a synthesised interior (category) node."""

    token: str                                              # last path segment; "" for the root
    children: dict[str, "CommandNode"] = field(default_factory=dict)
    entry: dict | None = None

    @property
    def is_leaf(self) -> bool:
        return self.entry is not None


class CommandTree:
    def __init__(self, commands: dict):
        self.root = CommandNode("")
        for path, entry in commands.items():
            node = self.root
            for seg in path.split():
                node = node.children.setdefault(seg, CommandNode(seg))
            node.entry = entry

    def node_at(self, tokens) -> CommandNode | None:
        """The node reached by walking ``tokens`` from the root, or ``None`` if any is unknown."""
        node = self.root
        for tok in tokens:
            node = node.children.get(tok)
            if node is None:
                return None
        return node

    def children_of(self, tokens) -> list[CommandNode]:
        """Completion candidates below ``tokens`` (alphabetical), or ``[]`` for an unknown path."""
        node = self.node_at(tokens)
        if node is None:
            return []
        return [node.children[k] for k in sorted(node.children)]

    def entry_of(self, tokens) -> dict | None:
        """The leaf metadata at ``tokens`` (``None`` for an interior or unknown path)."""
        node = self.node_at(tokens)
        return node.entry if node else None


def parse_option(line: str):
    """Split a rendered option string into ``(flags, metavar, description)``.

    ``"-k, --key <hex> key, 6 hex bytes"`` -> ``(["-k", "--key"], "<hex>", "key, 6 hex bytes")``
    ``"--blk <dec> block number"``         -> ``(["--blk"], "<dec>", "block number")``
    ``"-a input key type is key A (def)"`` -> ``(["-a"], "", "input key type is key A (def)")``

    The leading run of flag tokens is delimited by trailing commas (``-k,`` promises another flag);
    a following ``<...>`` token is the metavar; the rest is the description.
    """
    toks = line.split()
    flags: list[str] = []
    i, more = 0, True
    while i < len(toks) and more:
        raw = toks[i]
        core = raw[:-1] if raw.endswith(",") else raw
        if not _FLAG.match(core):
            break
        flags.append(core)
        more = raw.endswith(",")                       # a trailing comma promises another flag
        i += 1
    meta = ""
    if i < len(toks) and toks[i].startswith("<"):
        meta = toks[i]
        i += 1
    return flags, meta, " ".join(toks[i:])


def load_command_tree(path=None):
    """Return ``(CommandTree, metadata)``. Resolution order: an explicit ``path`` (``--commands``),
    then ``$PM3_COMMANDS_JSON``, then the discovered repo's ``doc/commands.json`` (see
    :func:`find_repo_root`), then ``./doc/commands.json``. Raises :class:`FileNotFoundError` (with
    the regen hint) when nothing is found."""
    candidates = []
    if path:
        candidates.append(Path(path))
    elif os.environ.get(_ENV_VAR):
        candidates.append(Path(os.environ[_ENV_VAR]))
    else:
        root = find_repo_root()
        if root is not None:
            candidates.append(root / "doc" / "commands.json")
        candidates.append(Path("doc/commands.json"))
    for cand in candidates:
        p = cand.expanduser()
        if p.is_file():
            data = json.loads(p.read_text())
            return CommandTree(data["commands"]), data.get("metadata", {})
    raise FileNotFoundError(
        "no Proxmark3 commands.json found (looked in: "
        + ", ".join(str(c) for c in candidates) + ").\n"
        "  run this from inside a Proxmark3 checkout, pass --commands PATH, set $PM3_REPO, or\n"
        "  regenerate it:  proxmark3 --fulltext | client/pyscripts/pm3_help2json.py - -"
    )