API Reference¶

class TKeySign(TKey):
    """Client for communicating with the TKey signer application.

    This class implements public key retrieval and signing as defined in the
    [tkey-pq-device-signer protocol](https://github.com/tillitis/tkey-pq-device-signer)
    but is also compatible with the
    [tkey-device-signer protocol](https://github.com/tillitis/tkey-device-signer).
    """

    def __init__(
        self,
        app: SignApp,
        device: str | None = None,
        secret: str | None = None,
        *,
        require_firmware_mode: bool = False,
    ) -> None:
        """Initialize the TKey signing client.

        If the TKey device is in firmware mode, this will automatically load the
        application binary. If the device is already running an application (and
        require_firmware_mode is not set), verifies that the running application
        matches the expected name and version.

        Args:
            app: The SignApp configuration containing the binary and metadata.
            device: Optional serial port path (e.g., `/dev/ttyACM0`). If None,
                the port is auto-detected.
            secret: Optional User Supplied Secret (passphrase) used as a seed
                for key derivation.
            require_firmware_mode: If True, fail if the device is not in firmware
                mode instead of attempting to use an already running application.

        Raises:
            TKeyNotFoundError: If the TKey device cannot be found.
            TKeyNotInFirmwareModeError: when require_firmware_mode is set and the
                device is not in firmware mode.
            TKeyUnexpectedAppError: If loading the application fails or the device is
                running a mismatched application.
            TKeyError: For other connection or initialization failures.
        """
        super().__init__(device)
        self.key_size = app.key_size
        self.sig_size = app.sig_size
        self.name = app.name

        try:
            if not self.load_app(app.binary, secret):
                if require_firmware_mode:
                    raise TKeyNotInFirmwareModeError(
                        "TKey is not in firmware mode but require_firmware_mode was set"
                    )

                # TKey is not in firmware mode: Query application name and version
                rx = self.send(SignCmd.GET_NAMEVERSION)
                name = (
                    rx[2:6].decode("ascii").rstrip(),
                    rx[6:10].decode("ascii").rstrip(),
                )
                ver = int.from_bytes(rx[10:14], byteorder="little")
                if name == app.name and ver == app.version:
                    return  # Signer application is already loaded

                raise TKeyUnexpectedAppError(
                    f"TKey is running an unknown application {name, ver}, "
                    f"expected {app.name, app.version}"
                )
        except TKeyError:
            self.disconnect()
            raise

    def get_pubkey(self) -> bytes:
        """Retrieve the public key bytes from the TKey device.

        Returns:
            The raw public key bytes.

        Raises:
            TKeyIOError: If reading from the serial port fails.
            TKeyProtocolError: If there is a framing or protocol mismatch.
        """
        pubkey = bytearray(self.key_size)

        # Issue command, read first frame
        rx = self.send(SignCmd.GET_PUBKEY)
        offset = 0

        while offset < self.key_size:
            chunk_size = min(self.key_size - offset, 127)
            pubkey[offset : offset + chunk_size] = rx[2 : 2 + chunk_size]
            offset += chunk_size
            if offset < self.key_size:
                rx = self.recv_response(SignCmd.GET_PUBKEY)

        return bytes(pubkey)

    def sign(self, message: SignableMessage, pub_key: bytes | None = None) -> bytes:
        """Sign a payload.

        Sends payload to device and retrieves the signature.

        For ML-DSA, the FIPS 204 external mu is computed using the message
        and public key: the mu is sent to device instead of payload.
        Streaming messages of arbitrary size are supported for ML-DSA.
        For Ed25519, the message (or stream) must not exceed 4096 bytes.

        Note:
            This method blocks and waits (up to 60 seconds) for the user to touch
            the physical TKey device when it flashes.

        Args:
            message: The raw bytes, binary reader (file-like object), or chunk iterable
                to sign. For Ed25519, the total message size cannot exceed 4096B.
                This limitation does not apply to ML-DSA as FIPS 204 external mu
                is used.
            pub_key: The public key bytes (only needed for ML-DSA). If not provided,
                key is retrieved from device.

        Returns:
            The generated signature as raw bytes.

        Raises:
            TKeyError: If the device returns a bad status during signing.
            TKeyIOError: If writing or reading from the serial port fails.
            TKeyProtocolError: If there is a framing or protocol mismatch.
        """
        if self.name == _MLDSA_APP_NAME:
            # Compute FIPS 204 external mu
            if pub_key is None:
                pub_key = self.get_pubkey()
            tr = hashlib.shake_256(pub_key).digest(64)
            shake = hashlib.shake_256(tr + b"\x00\x00")
            for chunk in _iter_chunks(message):
                shake.update(chunk)
            payload = shake.digest(64)
        else:
            payload = _read_bounded(message, MAX_PAYLOAD_SIZE)

        # Set size
        self.send(SignCmd.SET_SIZE, len(payload).to_bytes(4, byteorder="little"))

        # Load data in chunks
        chunk_size = 127
        offset = 0
        while offset < len(payload):
            chunk = payload[offset : offset + chunk_size]
            rx = self.send(SignCmd.LOAD_DATA, chunk)
            if rx[2] != 0:
                raise TKeyError(f"LoadData chunk NOK status: {rx[2]}")
            offset += chunk_size

        # Trigger signing (blocks waiting for touch) and read first frame
        rx = self.send(SignCmd.GET_SIG, timeout=60)

        # Read remaining frames
        signature = bytearray(self.sig_size)
        offset = 0

        while offset < self.sig_size:
            if rx[2] != 0:
                raise TKeyError(f"GetSig chunk NOK status: {rx[2]}")

            chunk_size = min(self.sig_size - offset, 126)
            signature[offset : offset + chunk_size] = rx[3 : 3 + chunk_size]
            offset += chunk_size

            if offset < self.sig_size:
                rx = self.recv_response(SignCmd.GET_SIG)

        return bytes(signature)

@dataclass
class SignApp:
    """Configuration and binary data for the TKey device signer application.

    Attributes:
        binary: The raw bytes of the device application binary.
        version: The version number of the device application.
        name: Device application name tuple.
        sig_size: The size of the generated signature in bytes.
        key_size: The size of the public key in bytes.
    """

    binary: bytes
    version: int
    name: tuple[str, str]
    sig_size: int
    key_size: int

    @property
    def digest(self) -> str:
        """Return the SHA-512 hex digest of the application binary."""
        return hashlib.sha512(self.binary).hexdigest()

    @classmethod
    def _find_binary(
        cls, version: int | None, digest: str | None, bins: list[tuple[str, int]]
    ) -> tuple[bytes, int]:
        resources_dir = importlib.resources.files("keylet.resources")
        matches = []

        # Scan registered binaries
        for filename, file_ver in bins:
            # Filter by version if requested
            if version is not None and file_ver != version:
                continue

            binary = resources_dir.joinpath(filename).read_bytes()
            file_digest = hashlib.sha512(binary).hexdigest()

            # Filter by digest if requested
            if digest is not None and not file_digest.startswith(digest.lower()):
                continue

            matches.append((binary, file_ver))

            if digest is None and version is None:
                # First binary is the default one
                break

        if not matches:
            raise ValueError(
                f"No device binary found matching: version={version}, digest={digest}"
            )

        if len(matches) > 1:
            raise ValueError(
                f"Multiple device binaries found matching: version={version}, "
                f"digest={digest}."
            )

        return matches[0]

    @classmethod
    def load_mldsa(
        cls, version: int | None = None, digest: str | None = None
    ) -> SignApp:
        """Load a ML-DSA signer application from package resources.

        If a digest (or prefix) is provided, it returns the binary matching the
        digest. If a version is provided, it filters by version. If neither is
        provided, current default binary is loaded.

        TKey key derivation depends on the application binary, so users who want a
        specific key must provide the binary digest.

        Args:
            version: The version of the signer application to load.
            digest: A SHA-512 hex digest (or prefix) of the target binary.

        Returns:
            An instance of SignApp configured with the loaded binary.

        Raises:
            ValueError: If no binary matches the criteria, or if the search
                is ambiguous (matches multiple binaries).
        """

        binary, version = cls._find_binary(version, digest, _EMBEDDED_MLDSA_BINS)
        return cls(binary, version, _MLDSA_APP_NAME, 2420, 1312)

    @classmethod
    def load_ed25519(
        cls, version: int | None = None, digest: str | None = None
    ) -> SignApp:
        """Load a Ed25519 signer application from package resources.

        If a digest (or prefix) is provided, it returns the binary matching the
        digest. If a version is provided, it filters by version. If neither is
        provided, current default binary is loaded.

        TKey key derivation depends on the application binary, so users who want a
        specific key must provide the binary digest.

        Warning:
            When Ed25519 is used, there is a 4096B size limit to signing payloads.

        Args:
            version: The version of the signer application to load.
            digest: A SHA-512 hex digest (or prefix) of the target binary.

        Returns:
            An instance of SignApp configured with the loaded binary.

        Raises:
            ValueError: If no binary matches the criteria, or if the search
                is ambiguous (matches multiple binaries).
        """
        binary, version = cls._find_binary(version, digest, _EMBEDDED_ED25519_BINS)
        return cls(binary, version, _ED25519_APP_NAME, 64, 32)