ledger-handshake/tests/application_client/handshake_command_sender.py

"""APDU client for the Handshake Ledger app.

Speaks the protocol the app actually implements: `SIGN_TX` is a state machine
selected by P1 (BEGIN / ADD_INPUT / ADD_OUTPUT / REVIEW / SIGN_INPUT), not the
chunked single-shot flow the Ledger boilerplate uses. See
`src/handlers/sign_tx.rs` and `scripts/sign-tx`.
"""

from contextlib import contextmanager
from enum import IntEnum
from typing import Generator, List, Optional

from ragger.backend.interface import BackendInterface, RAPDU
from ragger.bip import pack_derivation_path

CLA: int = 0xE0

#: Data bytes per APDU. ADD_OUTPUT payloads above this are chunked with P2_MORE.
MAX_APDU_DATA: int = 240


class InsType(IntEnum):
    GET_VERSION = 0x03
    GET_APP_NAME = 0x04
    GET_PUBLIC_KEY = 0x05
    SIGN_TX = 0x06
    SIGN_MESSAGE = 0x07


class P1(IntEnum):
    NONE = 0x00
    #: GET_PUBLIC_KEY: show the address on-device before replying.
    CONFIRM = 0x01
    #: SIGN_TX sub-ops.
    SIGN_BEGIN = 0x00
    SIGN_ADD_INPUT = 0x01
    SIGN_ADD_OUTPUT = 0x02
    SIGN_REVIEW = 0x03
    SIGN_INPUT = 0x04


class P2(IntEnum):
    LAST = 0x00
    MORE = 0x80


class Errors(IntEnum):
    """Mirrors `AppSW` in src/main.rs."""

    SW_DENY = 0x6985
    SW_WRONG_P1P2 = 0x6A86
    SW_WRONG_STATE = 0x6B00
    SW_INS_NOT_SUPPORTED = 0x6D00
    SW_CLA_NOT_SUPPORTED = 0x6E00
    SW_COMM_ERROR = 0x6F00
    SW_TX_DISPLAY_FAIL = 0xB001
    SW_ADDR_DISPLAY_FAIL = 0xB002
    SW_TX_WRONG_LENGTH = 0xB004
    SW_TX_PARSING_FAIL = 0xB005
    SW_TX_HASH_FAIL = 0xB006
    SW_TX_SIGN_FAIL = 0xB008
    SW_KEY_DERIVE_FAIL = 0xB009
    SW_VERSION_PARSING_FAIL = 0xB00A
    SW_BAD_DERIVATION_PATH = 0xB00B
    SW_SWAP_FAIL = 0xC000


def split_message(message: bytes, max_size: int) -> List[bytes]:
    return [message[x:x + max_size] for x in range(0, len(message), max_size)]


class HandshakeCommandSender:
    def __init__(self, backend: BackendInterface) -> None:
        self.backend = backend

    # ── Simple queries ─────────────────────────────────────────────────────

    def get_version(self) -> RAPDU:
        return self.backend.exchange(cla=CLA, ins=InsType.GET_VERSION,
                                     p1=P1.NONE, p2=P2.LAST, data=b"")

    def get_app_name(self) -> RAPDU:
        return self.backend.exchange(cla=CLA, ins=InsType.GET_APP_NAME,
                                     p1=P1.NONE, p2=P2.LAST, data=b"")

    def get_public_key(self, path: str) -> RAPDU:
        return self.backend.exchange(cla=CLA, ins=InsType.GET_PUBLIC_KEY,
                                     p1=P1.NONE, p2=P2.LAST,
                                     data=pack_derivation_path(path))

    @contextmanager
    def get_public_key_with_confirmation(self, path: str) -> Generator[None, None, None]:
        with self.backend.exchange_async(cla=CLA, ins=InsType.GET_PUBLIC_KEY,
                                         p1=P1.CONFIRM, p2=P2.LAST,
                                         data=pack_derivation_path(path)) as response:
            yield response

    # ── SIGN_TX state machine ──────────────────────────────────────────────

    def sign_tx_begin(self, payload: bytes) -> RAPDU:
        return self.backend.exchange(cla=CLA, ins=InsType.SIGN_TX,
                                     p1=P1.SIGN_BEGIN, p2=P2.LAST, data=payload)

    def sign_tx_add_input(self, payload: bytes) -> RAPDU:
        return self.backend.exchange(cla=CLA, ins=InsType.SIGN_TX,
                                     p1=P1.SIGN_ADD_INPUT, p2=P2.LAST, data=payload)

    def sign_tx_add_output(self, payload: bytes) -> RAPDU:
        """Send one output, chunking it across APDUs if needed."""
        chunks = split_message(payload, MAX_APDU_DATA) or [b""]
        rapdu = None
        for i, chunk in enumerate(chunks):
            p2 = P2.LAST if i == len(chunks) - 1 else P2.MORE
            rapdu = self.backend.exchange(cla=CLA, ins=InsType.SIGN_TX,
                                          p1=P1.SIGN_ADD_OUTPUT, p2=p2, data=chunk)
        return rapdu

    def sign_tx_review_sync(self) -> RAPDU:
        """REVIEW without a navigator: only for cases the device rejects
        before it puts anything on screen."""
        return self.backend.exchange(cla=CLA, ins=InsType.SIGN_TX,
                                     p1=P1.SIGN_REVIEW, p2=P2.LAST, data=b"")

    @contextmanager
    def sign_tx_review(self) -> Generator[None, None, None]:
        """REVIEW: blocks on the user, so the caller drives the navigator."""
        with self.backend.exchange_async(cla=CLA, ins=InsType.SIGN_TX,
                                         p1=P1.SIGN_REVIEW, p2=P2.LAST,
                                         data=b"") as response:
            yield response

    def sign_tx_input(self, index: int) -> RAPDU:
        return self.backend.exchange(cla=CLA, ins=InsType.SIGN_TX,
                                     p1=P1.SIGN_INPUT, p2=P2.LAST,
                                     data=index.to_bytes(4, "big"))

    # ── SIGN_MESSAGE ───────────────────────────────────────────────────────

    @contextmanager
    def sign_message(self, path: str, message: bytes) -> Generator[None, None, None]:
        payload = (pack_derivation_path(path)
                   + len(message).to_bytes(2, "little")
                   + message)
        with self.backend.exchange_async(cla=CLA, ins=InsType.SIGN_MESSAGE,
                                         p1=P1.NONE, p2=P2.LAST,
                                         data=payload) as response:
            yield response

    def get_async_response(self) -> Optional[RAPDU]:
        return self.backend.last_async_response