from enum import IntEnum from typing import List, Mapping from collections import deque from hashlib import sha256 from .common import ByteStreamParser, sha256, write_varint from .merkle import MerkleTree, element_hash class ClientCommandCode(IntEnum): YIELD = 0x10 GET_PREIMAGE = 0x40 GET_MERKLE_LEAF_PROOF = 0x41 GET_MERKLE_LEAF_INDEX = 0x42 GET_MORE_ELEMENTS = 0xA0 CCMD_YIELD_MUSIG_PUBNONCE_TAG = 0xFFFFFFFF CCMD_YIELD_MUSIG_PARTIALSIGNATURE_TAG = 0xFFFFFFFE class ClientCommand: def execute(self, request: bytes) -> bytes: raise NotImplementedError("Subclasses should implement this method.") @property def code(self) -> int: raise NotImplementedError("Subclasses should implement this method.") class YieldCommand(ClientCommand): def __init__(self, results: List[bytes]): self.results = results @property def code(self) -> int: return ClientCommandCode.YIELD def execute(self, request: bytes) -> bytes: self.results.append(request[1:]) # only skip the first byte (command code) return b"" class GetPreimageCommand(ClientCommand): def __init__(self, known_preimages: Mapping[bytes, bytes], queue: "deque[bytes]"): self.queue = queue self.known_preimages = known_preimages @property def code(self) -> int: return ClientCommandCode.GET_PREIMAGE def execute(self, request: bytes) -> bytes: req = ByteStreamParser(request[1:]) if req.read_bytes(1) != b'\0': raise RuntimeError(f"Unsupported request: the first byte should be 0") req_hash = req.read_bytes(32) req.assert_empty() if req_hash in self.known_preimages: known_preimage = self.known_preimages[req_hash] preimage_len_out = write_varint(len(known_preimage)) # We can send at most 255 - len(preimage_len_out) - 1 bytes in a single message; # the rest will be stored for GET_MORE_ELEMENTS max_payload_size = 255 - len(preimage_len_out) - 1 payload_size = min(max_payload_size, len(known_preimage)) if payload_size < len(known_preimage): # split into list of length-1 bytes elements extra_elements = [ known_preimage[i: i + 1] for i in range(payload_size, len(known_preimage)) ] # add to the queue any remaining extra bytes self.queue.extend(extra_elements) return ( preimage_len_out + payload_size.to_bytes(1, byteorder="big") + known_preimage[:payload_size] ) # not found raise RuntimeError(f"Requested unknown preimage for: {req_hash.hex()}") class GetMerkleLeafProofCommand(ClientCommand): def __init__(self, known_trees: Mapping[bytes, MerkleTree], queue: "deque[bytes]"): self.queue = queue self.known_trees = known_trees @property def code(self) -> int: return ClientCommandCode.GET_MERKLE_LEAF_PROOF def execute(self, request: bytes) -> bytes: req = ByteStreamParser(request[1:]) root = req.read_bytes(32) tree_size = req.read_varint() leaf_index = req.read_varint() req.assert_empty() if not root in self.known_trees: raise ValueError(f"Unknown Merkle root: {root.hex()}.") mt: MerkleTree = self.known_trees[root] if leaf_index >= tree_size or len(mt) != tree_size: raise ValueError(f"Invalid index or tree size.") if len(self.queue) != 0: raise RuntimeError( "This command should not execute when the queue is not empty." ) proof = mt.prove_leaf(leaf_index) # Compute how many elements we can fit in 255 - 32 - 1 - 1 = 221 bytes n_response_elements = min((255 - 32 - 1 - 1) // 32, len(proof)) n_leftover_elements = len(proof) - n_response_elements # Add to the queue any proof elements that do not fit the response if (n_leftover_elements > 0): self.queue.extend(proof[-n_leftover_elements:]) return b"".join( [ mt.get(leaf_index), len(proof).to_bytes(1, byteorder="big"), n_response_elements.to_bytes(1, byteorder="big"), *proof[:n_response_elements], ] ) class GetMerkleLeafIndexCommand(ClientCommand): def __init__(self, known_trees: Mapping[bytes, MerkleTree]): self.known_trees = known_trees @property def code(self) -> int: return ClientCommandCode.GET_MERKLE_LEAF_INDEX def execute(self, request: bytes) -> bytes: req = ByteStreamParser(request[1:]) root = req.read_bytes(32) leaf_hash = req.read_bytes(32) req.assert_empty() if root not in self.known_trees: raise ValueError(f"Unknown Merkle root: {root.hex()}.") try: leaf_index = self.known_trees[root].leaf_index(leaf_hash) found = 1 except ValueError: leaf_index = 0 found = 0 return found.to_bytes(1, byteorder="big") + write_varint(leaf_index) class GetMoreElementsCommand(ClientCommand): def __init__(self, queue: "deque[bytes]"): self.queue = queue @property def code(self) -> int: return ClientCommandCode.GET_MORE_ELEMENTS def execute(self, request: bytes) -> bytes: if len(request) != 1: raise ValueError("Wrong request length.") if len(self.queue) == 0: raise ValueError("No elements to get.") element_len = len(self.queue[0]) if any(len(el) != element_len for el in self.queue): raise ValueError( "The queue contains elements of different byte length, which is not expected." ) # pop from the queue, keeping the total response length at most 255 response_elements = bytearray() n_added_elements = 0 while len(self.queue) > 0 and len(response_elements) + element_len <= 253: response_elements.extend(self.queue.popleft()) n_added_elements += 1 return b"".join( [ n_added_elements.to_bytes(1, byteorder="big"), element_len.to_bytes(1, byteorder="big"), bytes(response_elements), ] ) class ClientCommandInterpreter: """Interpreter for the client-side commands. This class keeps has methods to keep track of: - known preimages - known Merkle trees from lists of elements Moreover, it containes the state that is relevant for the interpreted client side commands: - a queue of bytes that contains any bytes that could not fit in a response from the GET_PREIMAGE client command (when a preimage is too long to fit in a single message) or the GET_MERKLE_LEAF_PROOF command (which returns a Merkle proof, which might be too long to fit in a single message). The data in the queue is returned in one (or more) successive GET_MORE_ELEMENTS commands from the hardware wallet. Finally, it keeps track of the yielded values (that is, the values sent from the hardware wallet with a YIELD client command). Attributes ---------- yielded: list[bytes] A list of all the value sent by the Hardware Wallet with a YIELD client command during thw processing of an APDU. """ def __init__(self): self.known_preimages: Mapping[bytes, bytes] = {} self.known_trees: Mapping[bytes, MerkleTree] = {} self.yielded: List[bytes] = [] queue = deque() commands = [ YieldCommand(self.yielded), GetPreimageCommand(self.known_preimages, queue), GetMerkleLeafIndexCommand(self.known_trees), GetMerkleLeafProofCommand(self.known_trees, queue), GetMoreElementsCommand(queue), ] self.commands = {cmd.code: cmd for cmd in commands} def execute(self, hw_response: bytes) -> bytes: """Interprets the client command requested by the hardware wallet, returning the appropriate response and updating the client interpreter's internal state if needed. Parameters ---------- hw_response : bytes The data content of the SW_INTERRUPTED_EXECUTION sent by the hardware wallet. Returns ------- bytes The result of the execution of the appropriate client side command, containing the response to be sent via INS_CONTINUE. """ if len(hw_response) == 0: raise RuntimeError( "Unexpected empty SW_INTERRUPTED_EXECUTION response from hardware wallet." ) cmd_code = hw_response[0] if cmd_code not in self.commands: raise RuntimeError( "Unexpected command code: 0x{:02X}".format(cmd_code) ) return self.commands[cmd_code].execute(hw_response) def add_known_preimage(self, element: bytes) -> None: """Adds a preimage to the list of known preimages. The client must respond with `element` when a GET_PREIMAGE command is sent with `sha256(element)` in its request. Parameters ---------- element : bytes An array of bytes whose preimage must be known to the client during an APDU execution. """ self.known_preimages[sha256(element)] = element def add_known_list(self, elements: List[bytes]) -> None: """Adds a known Merkleized list. Builds the Merkle tree of `elements`, and adds it to the Merkle trees known to the client (mapped by Merkle root `mt_root`). moreover, adds all the leafs (after adding the b'\0' prefix) to the list of known preimages. If `el` is one of `elements`, the client must respond with b'\0' + `el` when a GET_PREIMAGE client command is sent with `sha256(b'\0' + el)`. Moreover, the commands GET_MERKLE_LEAF_INDEX and GET_MERKLE_LEAF_PROOF must correctly answer queries relative to the Merkle whose root is `mt_root`. Parameters ---------- elements : List[bytes] A list of `bytes` corresponding to the leafs of the Merkle tree. """ for el in elements: self.add_known_preimage(b"\x00" + el) mt = MerkleTree(element_hash(el) for el in elements) self.known_trees[mt.root] = mt def add_known_mapping(self, mapping: Mapping[bytes, bytes]) -> None: """Adds the Merkle trees of keys, and the Merkle tree of values (ordered by key) of a mapping of bytes to bytes. Adds the Merkle tree of the list of keys, and the Merkle tree of the list of corresponding values, with the same semantics as the `add_known_list` applied separately to the two lists. Parameters ---------- mapping : Mapping[bytes, bytes] A mapping whose keys and values are `bytes`. """ items_sorted = list(sorted(mapping.items())) keys = [i[0] for i in items_sorted] values = [i[1] for i in items_sorted] self.add_known_list(keys) self.add_known_list(values)