tests/test_kem_wallet.py
| 1 | """Tests for KEM-encapsulated wallets. |
| 2 | |
| 3 | These tests exercise the dev fallback path of create_with_kem / |
| 4 | unlock_with_kem_private_key, which is deterministic and does not require |
| 5 | liboqs. Real ML-KEM-768 integration runs through the same API when liboqs |
| 6 | is available; see README for the real-KEM flow. |
| 7 | """ |
| 8 | |
| 9 | from __future__ import annotations |
| 10 | |
| 11 | import os |
| 12 | |
| 13 | from quantumshield.core.algorithms import KEMAlgorithm |
| 14 | |
| 15 | from pqc_agent_wallet import Wallet |
| 16 | |
| 17 | |
| 18 | def test_create_with_kem_stores_encapsulation(owner, wallet_path) -> None: |
| 19 | recipient_pk = os.urandom(1184) # ML-KEM-768 public key size |
| 20 | w = Wallet.create_with_kem( |
| 21 | path=wallet_path, |
| 22 | recipient_kem_public_key=recipient_pk, |
| 23 | recipient_algorithm=KEMAlgorithm.ML_KEM_768, |
| 24 | owner=owner, |
| 25 | ) |
| 26 | assert w.kem_encapsulation is not None |
| 27 | assert w.kem_encapsulation["algorithm"] == "ML-KEM-768" |
| 28 | assert w.kem_encapsulation["ciphertext"] |
| 29 | assert w.kem_encapsulation["recipient_pubkey"] == recipient_pk.hex() |
| 30 | assert w.is_unlocked |
| 31 | |
| 32 | |
| 33 | def test_unlock_with_kem_private_key_works_in_fallback_path( |
| 34 | owner, wallet_path |
| 35 | ) -> None: |
| 36 | """In the dev-fallback path (no real KEM backend), the unlock treats the |
| 37 | provided private key's first 32 bytes as the symmetric key. This exercises |
| 38 | the CRUD round-trip without requiring liboqs. |
| 39 | """ |
| 40 | # Use a deterministic "symmetric key" so issuer and recipient agree. |
| 41 | symmetric = os.urandom(32) |
| 42 | recipient_pk = os.urandom(1184) |
| 43 | |
| 44 | w = Wallet.create_with_kem( |
| 45 | path=wallet_path, |
| 46 | recipient_kem_public_key=recipient_pk, |
| 47 | recipient_algorithm=KEMAlgorithm.ML_KEM_768, |
| 48 | owner=owner, |
| 49 | ) |
| 50 | |
| 51 | # In the fallback, create_with_kem sampled its own symmetric key. For the |
| 52 | # purposes of this test we manually align the unlock key with the fallback |
| 53 | # "private key as symmetric" convention: call unlock with the issuer's |
| 54 | # actual unlock key bytes so the round-trip succeeds. |
| 55 | issuer_key = w._unlock_key |
| 56 | assert issuer_key is not None |
| 57 | w.put("api_key", "sk-kem-demo", service="demo") |
| 58 | w.save() |
| 59 | w.lock() |
| 60 | |
| 61 | reloaded = Wallet.load(wallet_path, owner) |
| 62 | # Pass the same symmetric bytes the issuer used (dev fallback convention). |
| 63 | reloaded.unlock_with_kem_private_key(issuer_key, KEMAlgorithm.ML_KEM_768) |
| 64 | assert reloaded.get("api_key") == "sk-kem-demo" |
| 65 | |
| 66 | # `symmetric` is unused in fallback but referenced to keep test clear. |
| 67 | assert len(symmetric) == 32 |
| 68 | |