Encrypt remote transport with post-quantum session keys

browser_cli/auth.py

@@ -10,6 +10,9 @@ from pathlib import Path
 
 from cryptography.exceptions import InvalidSignature
 from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 from cryptography.hazmat.primitives.serialization import (
     Encoding,
     NoEncryption,
@@ -190,6 +193,7 @@ def verify(pub_hex: str, nonce: bytes, msg: dict, sig_hex: str, pq_shared_secret
 # ── Post-quantum key exchange (ML-KEM / Kyber) ────────────────────────────────
 
 PQ_KEX_ALG = "ML-KEM-768"
+PQ_TRANSPORT_ALG = "ML-KEM-768+ChaCha20Poly1305"
 
 
 def pq_kex_server_keypair():
@@ -221,6 +225,35 @@ def pq_kex_server_decapsulate(private_key, ciphertext_hex: str) -> bytes:
     return private_key.decapsulate(bytes.fromhex(ciphertext_hex))
 
 
+def _pq_transport_key(shared_secret: bytes, direction: str) -> bytes:
+    return HKDF(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=None,
+        info=f"browser-cli pq transport v1 {direction}".encode("ascii"),
+    ).derive(shared_secret)
+
+
+def pq_encrypt(shared_secret: bytes, direction: str, plaintext: bytes) -> dict:
+    """Encrypt an app-layer frame with a key derived from the ML-KEM secret."""
+    nonce = secrets.token_bytes(12)
+    key = _pq_transport_key(shared_secret, direction)
+    ciphertext = ChaCha20Poly1305(key).encrypt(nonce, plaintext, None)
+    return {"alg": PQ_TRANSPORT_ALG, "nonce": nonce.hex(), "ciphertext": ciphertext.hex()}
+
+
+def pq_decrypt(shared_secret: bytes, direction: str, envelope: dict) -> bytes:
+    """Decrypt an app-layer frame produced by pq_encrypt()."""
+    if not isinstance(envelope, dict) or envelope.get("alg") != PQ_TRANSPORT_ALG:
+        raise ValueError("unsupported encrypted transport envelope")
+    key = _pq_transport_key(shared_secret, direction)
+    return ChaCha20Poly1305(key).decrypt(
+        bytes.fromhex(str(envelope["nonce"])),
+        bytes.fromhex(str(envelope["ciphertext"])),
+        None,
+    )
+
+
 def new_nonce() -> str:
     return secrets.token_hex(32)