Add post-quantum remote auth key exchange

browser_cli/auth.py

@@ -1,4 +1,4 @@
-"""Ed25519 keypair management and challenge-response auth helpers."""
+"""Ed25519 keypair management, ML-KEM key exchange, and auth helpers."""
 import hashlib
 import json
 import os
@@ -152,34 +152,75 @@ def public_key_hex(key: Ed25519PrivateKey | AgentKey) -> str:
 # ── Canonical payload + sign/verify ───────────────────────────────────────────
 
 def canonical_payload(msg: dict) -> bytes:
-    """Deterministic JSON encoding of msg without auth fields."""
+    """Deterministic JSON encoding of msg without auth protocol fields."""
     return json.dumps(
-        {k: v for k, v in msg.items() if k not in {"pubkey", "sig"}},
+        {k: v for k, v in msg.items() if k not in {"pubkey", "sig", "pq_kex"}},
         sort_keys=True,
         separators=(",", ":"),
     ).encode("utf-8")
 
 
-def sign(key: Ed25519PrivateKey | AgentKey, nonce: bytes, msg: dict) -> bytes:
-    """Sign nonce + SHA256(canonical_payload(msg)) — works for both file keys and agent keys."""
+def _auth_message(nonce: bytes, msg: dict, pq_shared_secret: bytes | None = None) -> bytes:
+    """Bytes signed for auth; optionally binds a post-quantum KEX secret."""
     data = nonce + hashlib.sha256(canonical_payload(msg)).digest()
+    if pq_shared_secret is not None:
+        data += hashlib.sha256(b"browser-cli ml-kem-768 v1" + pq_shared_secret).digest()
+    return data
+
+
+def sign(key: Ed25519PrivateKey | AgentKey, nonce: bytes, msg: dict, pq_shared_secret: bytes | None = None) -> bytes:
+    """Sign nonce + payload hash, optionally bound to an ML-KEM shared secret."""
+    data = _auth_message(nonce, msg, pq_shared_secret)
     if isinstance(key, AgentKey):
         return agent_sign_raw(key, data)
     return key.sign(data)
 
 
-def verify(pub_hex: str, nonce: bytes, msg: dict, sig_hex: str) -> bool:
-    """Return True if sig_hex is a valid Ed25519 signature over the canonical payload."""
+def verify(pub_hex: str, nonce: bytes, msg: dict, sig_hex: str, pq_shared_secret: bytes | None = None) -> bool:
+    """Return True if sig_hex is a valid signature over the canonical payload/auth secret."""
     try:
         pub_bytes = bytes.fromhex(pub_hex)
         pub_key = Ed25519PublicKey.from_public_bytes(pub_bytes)
-        message = nonce + hashlib.sha256(canonical_payload(msg)).digest()
-        pub_key.verify(bytes.fromhex(sig_hex), message)
+        pub_key.verify(bytes.fromhex(sig_hex), _auth_message(nonce, msg, pq_shared_secret))
         return True
     except (InvalidSignature, ValueError):
         return False
 
 
+# ── Post-quantum key exchange (ML-KEM / Kyber) ────────────────────────────────
+
+PQ_KEX_ALG = "ML-KEM-768"
+
+
+def pq_kex_server_keypair():
+    """Return an ephemeral ML-KEM-768 private key and raw public key bytes.
+
+    Returns ``None`` when the installed cryptography/OpenSSL backend does not
+    support ML-KEM yet. The serve/client protocol treats this as graceful
+    downgrade instead of breaking local installs on older OpenSSL builds.
+    """
+    try:
+        from cryptography.hazmat.primitives.asymmetric import mlkem
+        priv = mlkem.MLKEM768PrivateKey.generate()
+        pub = priv.public_key().public_bytes_raw()
+        return priv, pub
+    except Exception:
+        return None
+
+
+def pq_kex_client_encapsulate(public_key_hex: str) -> tuple[str, bytes]:
+    """Encapsulate to a server ML-KEM public key. Returns (ciphertext_hex, secret)."""
+    from cryptography.hazmat.primitives.asymmetric import mlkem
+    pub = mlkem.MLKEM768PublicKey.from_public_bytes(bytes.fromhex(public_key_hex))
+    ciphertext, shared_secret = pub.encapsulate()
+    return ciphertext.hex(), shared_secret
+
+
+def pq_kex_server_decapsulate(private_key, ciphertext_hex: str) -> bytes:
+    """Decapsulate a client ML-KEM ciphertext and return the shared secret."""
+    return private_key.decapsulate(bytes.fromhex(ciphertext_hex))
+
+
 def new_nonce() -> str:
     return secrets.token_hex(32)