age plugin: align PQ host protocol with firmware; 1.2.11 packaging/description

MANIFEST.in

@@ -0,0 +1 @@
+include README.md

onlykey/age_plugin/__init__.py

@@ -11,11 +11,46 @@
 __version__ = "0.1.0"
 PLUGIN_NAME = "onlykey"
 
-# OnlyKey HID message types (from onlykey.client.Message)
-OKGETPUBKEY = 236
-OKDECRYPT = 240
-OKGENKEY = 230
-
-# OnlyKey reserved key slots for post-quantum
-SLOT_MLKEM = 133
-SLOT_XWING = 134
+# OnlyKey HID opcodes (values match onlykey.client.Message)
+OKGETPUBKEY = 236   # Message.OKGETPUBKEY
+OKDECRYPT = 240     # Message.OKDECRYPT
+OKSETPRIV = 239     # Message.OKSETPRIV — also generates an on-device key when
+                    # the key body is all 0xFF (firmware okcore.cpp set_private
+                    # -> okcrypto_generate_random_key)
+
+# ECC key slots that can hold the 32-byte post-quantum seed.
+#
+# A ML-KEM/X-Wing key is just a 32-byte seed stored in an ordinary ECC key
+# slot, and the algorithm is chosen by the key-type byte (buffer[6]) below —
+# not by the slot number. The slot is therefore caller-selectable, but only
+# across the USER key slots: firmware exposes 101-116 (16 slots) for user keys.
+# Slots 117-132 are RESERVED (e.g. 128 web-derivation, 129/130 HMAC, 131
+# backup, 132 derivation) and must never be used for PQ keys — writing there
+# would clobber internal device keys. Firmware getpubkey/decaps enforce this
+# with a `< 117` gate (okcrypto.cpp); the host mirrors it here.
+ECC_SLOT_MIN = 101
+ECC_SLOT_MAX = 116            # 16 user ECC key slots: 101-116 (117-132 reserved)
+DEFAULT_XWING_SLOT = 101
+DEFAULT_MLKEM_SLOT = 102
+
+
+def validate_ecc_slot(slot):
+    """Raise ValueError unless slot is a user ECC key slot (101-116).
+
+    117-132 are reserved by firmware and are intentionally rejected.
+    """
+    if not (ECC_SLOT_MIN <= int(slot) <= ECC_SLOT_MAX):
+        raise ValueError(
+            f"PQ key slot must be a user ECC slot {ECC_SLOT_MIN}-{ECC_SLOT_MAX} "
+            f"(117-132 are reserved), got {slot}"
+        )
+    return int(slot)
+
+# Firmware key-type identifiers (okcore.h: KEYTYPE_MLKEM768 / KEYTYPE_XWING).
+# Sent in the low nibble of buffer[6] so the device routes the operation.
+KEYTYPE_MLKEM768 = 5
+KEYTYPE_XWING = 6
+
+# An all-0xFF key body sent with OKSETPRIV tells the firmware to generate the
+# key on-device (gen_key trigger in okcore.cpp).
+GENERATE_ON_DEVICE = b"\xff" * 8

onlykey/age_plugin/cli.py

@@ -20,7 +20,9 @@
 import base64
 import hashlib
 
-from onlykey.age_plugin import __version__, PLUGIN_NAME, SLOT_XWING
+from onlykey.age_plugin import (
+    __version__, PLUGIN_NAME, DEFAULT_XWING_SLOT, validate_ecc_slot,
+)
 from onlykey.age_plugin.protocol import (
     Stanza, b64encode_no_pad, b64decode_no_pad,
     run_identity_v1, run_recipient_v1,
@@ -137,7 +139,7 @@ def recipient_fingerprint(pubkey: bytes) -> bytes:
     return hashlib.sha256(pubkey).digest()[:IDENTITY_FINGERPRINT_LEN]
 
 
-def encode_identity(slot: int = SLOT_XWING) -> str:
+def encode_identity(slot: int = DEFAULT_XWING_SLOT) -> str:
     """Encode an identity string. Contains just the slot number."""
     # Identity data: just the slot byte
     data = bytes([slot])
@@ -181,16 +183,16 @@ def decode_identity(identity: str) -> dict:
 
 
 
-def cmd_generate():
+def cmd_generate(slot: int = DEFAULT_XWING_SLOT):
     """Generate X-Wing keypair on OnlyKey and print recipient/identity."""
     from onlykey.age_plugin.onlykey_hid import OnlyKeyPQ
 
-    print("Generating X-Wing keypair on OnlyKey...", file=sys.stderr)
+    print(f"Generating X-Wing keypair on OnlyKey (ECC slot {slot})...", file=sys.stderr)
     dev = OnlyKeyPQ()
-    pk = dev.xwing_keygen()
+    pk = dev.xwing_keygen(slot)
 
     recipient = encode_recipient(pk)
-    identity = encode_identity(SLOT_XWING)
+    identity = encode_identity(slot)
 
     print("# X-Wing public key (produces native age mlkem768x25519 stanzas)", file=sys.stderr)
     print(f"# Recipient: {recipient}", file=sys.stderr)
@@ -202,19 +204,19 @@ def cmd_generate():
     print(identity)
 
 
-def cmd_recipient():
-    """Print the current X-Wing public key as an age recipient."""
+def cmd_recipient(slot: int = DEFAULT_XWING_SLOT):
+    """Print the X-Wing public key in the given ECC slot as an age recipient."""
     from onlykey.age_plugin.onlykey_hid import OnlyKeyPQ
 
     dev = OnlyKeyPQ()
-    pk = dev.xwing_getpubkey()
+    pk = dev.xwing_getpubkey(slot)
     print(encode_recipient(pk))
 
 
-def cmd_identity():
+def cmd_identity(slot: int = DEFAULT_XWING_SLOT):
     """Print an identity file for use with age -i."""
-    identity = encode_identity(SLOT_XWING)
-    print(f"# age-plugin-onlykey identity (X-Wing slot {SLOT_XWING})")
+    identity = encode_identity(slot)
+    print(f"# age-plugin-onlykey identity (X-Wing ECC slot {slot})")
     print(identity)
 
 
@@ -240,19 +242,27 @@ def unwrap_callback(identities, stanzas_per_file):
         return results
 
     dev = OnlyKeyPQ()
-    device_pubkey = dev.xwing_getpubkey()
-    if len(device_pubkey) != XWING_RECIPIENT_LEN:
-        raise ValueError(
-            f"OnlyKey returned an unexpected X-Wing public key length: {len(device_pubkey)}"
-        )
 
-    device_fingerprint = recipient_fingerprint(device_pubkey)
+    # The identity carries the ECC slot the key lives in; any of the 32 ECC
+    # slots (101-132) is valid. Query that slot's public key and match it.
     matching_identity = None
+    matching_slot = None
     for identity in parsed_identities:
-        if identity["slot"] != SLOT_XWING:
+        try:
+            slot = validate_ecc_slot(identity["slot"])
+        except ValueError:
             continue
-        if identity["fingerprint"] is None or identity["fingerprint"] == device_fingerprint:
+        try:
+            device_pubkey = dev.xwing_getpubkey(slot)
+        except Exception as exc:
+            print(f"Could not read X-Wing key in slot {slot}: {exc}", file=sys.stderr)
+            continue
+        if len(device_pubkey) != XWING_RECIPIENT_LEN:
+            continue
+        if (identity["fingerprint"] is None
+                or identity["fingerprint"] == recipient_fingerprint(device_pubkey)):
             matching_identity = identity
+            matching_slot = slot
             break
 
     if matching_identity is None:
@@ -293,8 +303,8 @@ def unwrap_callback(identities, stanzas_per_file):
                 )
 
 
-            # Send ciphertext to OnlyKey for decapsulation
-            ss = dev.xwing_decaps(enc)
+            # Send ciphertext to OnlyKey for decapsulation (in the matched slot)
+            ss = dev.xwing_decaps(enc, slot=matching_slot)
 
             # Use shared secret to decrypt the file key via HPKE
             try:
@@ -360,13 +370,26 @@ def main():
                 print(f"Unknown state machine: {state_machine}", file=sys.stderr)
                 sys.exit(1)
 
+    # Optional --slot N / --slot=N selects which ECC slot (101-132) holds the key.
+    slot = DEFAULT_XWING_SLOT
+    for i, arg in enumerate(args):
+        if arg == "--slot" and i + 1 < len(args):
+            slot = args[i + 1]
+        elif arg.startswith("--slot="):
+            slot = arg.split("=", 1)[1]
+    try:
+        slot = validate_ecc_slot(slot)
+    except ValueError as exc:
+        print(f"Error: {exc}", file=sys.stderr)
+        sys.exit(1)
+
     # Direct invocation modes
     if "--generate" in args or "-g" in args:
-        cmd_generate()
+        cmd_generate(slot)
     elif "--recipient" in args or "-r" in args:
-        cmd_recipient()
+        cmd_recipient(slot)
     elif "--identity" in args or "-i" in args:
-        cmd_identity()
+        cmd_identity(slot)
     elif "--version" in args or "-v" in args:
         print(f"age-plugin-onlykey {__version__}")
     elif "--help" in args or "-h" in args:
@@ -380,6 +403,7 @@ def main():
         print("  --generate    Generate X-Wing keypair on OnlyKey")
         print("  --recipient   Print recipient (public key) for encryption")
         print("  --identity    Print identity file for decryption")
+        print("  --slot N      User ECC slot 101-116 to use (default 101)")
         print("  --help        Show full help")
         print()
         print("Quick start:")

onlykey/age_plugin/onlykey_hid.py

@@ -12,7 +12,11 @@
 import time
 
 from onlykey.client import OnlyKey, Message
-from . import OKGETPUBKEY, OKDECRYPT, OKGENKEY, SLOT_MLKEM, SLOT_XWING
+from . import (
+    OKGETPUBKEY, OKDECRYPT, OKSETPRIV, GENERATE_ON_DEVICE,
+    DEFAULT_MLKEM_SLOT, DEFAULT_XWING_SLOT,
+    KEYTYPE_MLKEM768, KEYTYPE_XWING, validate_ecc_slot,
+)
 
 # Sizes
 XWING_PK_SIZE = 1216
@@ -51,11 +55,28 @@ def _connect(self):
                 time.sleep(0.5)
         raise RuntimeError("Could not connect to OnlyKey. Is it plugged in and unlocked?")
 
-    def _send_and_receive(self, msg_type, slot, payload=b"",
+    def _send_and_receive(self, msg_type, slot, payload=b"", key_type=None,
                           expected_size=0, timeout_ms=10000):
-        """Send a message and collect multi-packet response."""
-        self.ok.send_message(msg=msg_type, slot_id=slot, payload=payload)
+        """Send a SINGLE-packet request and collect the response.
 
+        Wire layout expected by firmware: buffer[5]=slot, buffer[6]=key type,
+        buffer[7:]=payload. ``key_type`` is placed in buffer[6] so the device
+        routes the request to the ML-KEM / X-Wing handler for the ECC slot.
+
+        This is only valid when the request payload fits in one 64-byte report
+        (keygen trigger, getpubkey). Large inputs that exceed one report — the
+        decapsulation ciphertext — must use the multi-packet send path; see
+        ``*_decaps`` below.
+        """
+        body = bytearray()
+        if key_type is not None:
+            body.append(key_type & 0x0F)   # firmware buffer[6]
+        body.extend(payload)
+        self.ok.send_message(msg=Message(msg_type), slot_id=slot, payload=body)
+        return self._read_response(expected_size=expected_size, timeout_ms=timeout_ms)
+
+    def _read_response(self, expected_size=0, timeout_ms=10000):
+        """Collect a (possibly multi-packet) response from the device."""
         result = bytearray()
         deadline = time.time() + timeout_ms / 1000
         while time.time() < deadline:
@@ -75,69 +96,83 @@ def _send_and_receive(self, msg_type, slot, payload=b"",
 
         return bytes(result[:expected_size] if expected_size else result)
 
-    def xwing_keygen(self):
-        """Generate X-Wing keypair. Returns 1216-byte public key."""
+    def _decaps(self, ciphertext, slot):
+        """Send a KEM ciphertext for on-device decapsulation; return 32-byte SS.
+
+        The ciphertext (1088 B for ML-KEM, 1120 B for X-Wing) is far larger than
+        one 64-byte HID report, so it is streamed with the multi-packet protocol
+        (``send_large_message2``) — the same path the OnlyKey CLI uses to send
+        RSA/ECDH ciphertext for OKDECRYPT. Each packet carries
+        [slot, 0xFF-or-final-length, <=57 bytes], which the firmware accumulates
+        into its large buffer. The device reads the key TYPE from the key stored
+        in ``slot`` (not from the packet), waits for a button press, then returns
+        the 32-byte shared secret.
+        """
+        print("Press OnlyKey button to confirm decryption...", file=sys.stderr)
+        self.ok.send_large_message2(
+            msg=Message(OKDECRYPT), payload=list(ciphertext), slot_id=slot,
+        )
+        return self._read_response(expected_size=32, timeout_ms=30000)
+
+    def xwing_keygen(self, slot=DEFAULT_XWING_SLOT):
+        """Generate an X-Wing keypair in the given ECC slot. Returns 1216-byte pubkey."""
+        slot = validate_ecc_slot(slot)
         print("Press OnlyKey button to confirm key generation...", file=sys.stderr)
         pk = self._send_and_receive(
-            OKGENKEY, SLOT_XWING,
+            OKSETPRIV, slot,
+            payload=GENERATE_ON_DEVICE, key_type=KEYTYPE_XWING,
             expected_size=XWING_PK_SIZE,
             timeout_ms=30000,
         )
         if len(pk) != XWING_PK_SIZE:
             raise RuntimeError(f"X-Wing keygen: got {len(pk)} bytes, expected {XWING_PK_SIZE}")
         return pk
 
-    def xwing_getpubkey(self):
-        """Get X-Wing public key. Returns 1216-byte public key."""
+    def xwing_getpubkey(self, slot=DEFAULT_XWING_SLOT):
+        """Get the X-Wing public key from the given ECC slot. Returns 1216-byte pubkey."""
+        slot = validate_ecc_slot(slot)
         pk = self._send_and_receive(
-            OKGETPUBKEY, SLOT_XWING,
+            OKGETPUBKEY, slot, key_type=KEYTYPE_XWING,
             expected_size=XWING_PK_SIZE,
             timeout_ms=10000,
         )
         if len(pk) != XWING_PK_SIZE:
             raise RuntimeError(f"X-Wing getpubkey: got {len(pk)} bytes, expected {XWING_PK_SIZE}")
         return pk
 
-    def xwing_decaps(self, ciphertext):
-        """X-Wing decapsulation. Returns 32-byte shared secret."""
+    def xwing_decaps(self, ciphertext, slot=DEFAULT_XWING_SLOT):
+        """X-Wing decapsulation in the given ECC slot. Returns 32-byte shared secret."""
+        slot = validate_ecc_slot(slot)
         if len(ciphertext) != XWING_CT_SIZE:
             raise ValueError(f"X-Wing CT must be {XWING_CT_SIZE} bytes, got {len(ciphertext)}")
-        print("Press OnlyKey button to confirm decryption...", file=sys.stderr)
-        ss = self._send_and_receive(
-            OKDECRYPT, SLOT_XWING,
-            payload=ciphertext,
-            expected_size=XWING_SS_SIZE,
-            timeout_ms=30000,
-        )
+        ss = self._decaps(ciphertext, slot)
         if len(ss) != XWING_SS_SIZE:
             raise RuntimeError(f"X-Wing decaps: got {len(ss)} bytes, expected {XWING_SS_SIZE}")
         return ss
 
-    def mlkem_keygen(self):
-        """Generate ML-KEM-768 keypair. Returns 1184-byte public key."""
+    def mlkem_keygen(self, slot=DEFAULT_MLKEM_SLOT):
+        """Generate an ML-KEM-768 keypair in the given ECC slot. Returns 1184-byte pubkey."""
+        slot = validate_ecc_slot(slot)
         print("Press OnlyKey button to confirm key generation...", file=sys.stderr)
         return self._send_and_receive(
-            OKGENKEY, SLOT_MLKEM,
+            OKSETPRIV, slot,
+            payload=GENERATE_ON_DEVICE, key_type=KEYTYPE_MLKEM768,
             expected_size=MLKEM_PK_SIZE,
             timeout_ms=30000,
         )
 
-    def mlkem_getpubkey(self):
-        """Get ML-KEM-768 public key. Returns 1184-byte public key."""
+    def mlkem_getpubkey(self, slot=DEFAULT_MLKEM_SLOT):
+        """Get the ML-KEM-768 public key from the given ECC slot. Returns 1184-byte pubkey."""
+        slot = validate_ecc_slot(slot)
         return self._send_and_receive(
-            OKGETPUBKEY, SLOT_MLKEM,
+            OKGETPUBKEY, slot, key_type=KEYTYPE_MLKEM768,
             expected_size=MLKEM_PK_SIZE,
             timeout_ms=10000,
         )
 
-    def mlkem_decaps(self, ciphertext):
-        """ML-KEM-768 decapsulation. Returns 32-byte shared secret."""
+    def mlkem_decaps(self, ciphertext, slot=DEFAULT_MLKEM_SLOT):
+        """ML-KEM-768 decapsulation in the given ECC slot. Returns 32-byte shared secret."""
+        slot = validate_ecc_slot(slot)
         if len(ciphertext) != MLKEM_CT_SIZE:
             raise ValueError(f"ML-KEM CT must be {MLKEM_CT_SIZE} bytes, got {len(ciphertext)}")
-        print("Press OnlyKey button to confirm decryption...", file=sys.stderr)
-        return self._send_and_receive(
-            OKDECRYPT, SLOT_MLKEM,
-            payload=ciphertext,
-            expected_size=32,
-            timeout_ms=30000,
-        )
+        return self._decaps(ciphertext, slot)

onlykey/client.py

@@ -11,7 +11,13 @@
 import os
 import codecs
 
-import hid
+try:
+    # On Linux the hidapi package ships a hidraw-backed module alongside the
+    # libusb-backed one. Prefer hidraw to avoid hid "open failed" races when the
+    # OnlyKey's HID interface was just used by another application (e.g. KeePassXC HMAC)
+    import hidraw as hid
+except ImportError:
+    import hid
 from aenum import Enum
 from sys import platform
 

setup.py

@@ -10,9 +10,10 @@
 
 setup(
     name='onlykey',
-    version='1.2.10',
+    version='1.2.11',
     description='OnlyKey client and command-line tool',
-    # long_description=long_description,
+    long_description=long_description,
+    long_description_content_type='text/markdown',
     url='https://github.com/trustcrypto/python-onlykey',
     author='CryptoTrust',
     author_email='admin@crp.to',