121 lines
4 KiB
Python
121 lines
4 KiB
Python
"""HPKE-Base X25519-HKDF-SHA256-ChaCha20-Poly1305 (RFC 9180).
|
|
|
|
Production E2E for NATS *content*. Routing fields stay in the clear.
|
|
The broker never gets the recipient private key.
|
|
"""
|
|
|
|
from __future__ import annotations # annotations
|
|
|
|
import os # eph nonce
|
|
from typing import Tuple # enc, ct
|
|
|
|
from cryptography.hazmat.primitives import hashes, hmac # HKDF pieces
|
|
from cryptography.hazmat.primitives.asymmetric.x25519 import ( # KEM
|
|
X25519PrivateKey,
|
|
X25519PublicKey,
|
|
)
|
|
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305 # AEAD
|
|
from cryptography.hazmat.primitives.serialization import ( # raw keys
|
|
Encoding,
|
|
NoEncryption,
|
|
PrivateFormat,
|
|
PublicFormat,
|
|
)
|
|
|
|
KEM_ID = 0x0020 # DHKEM(X25519, HKDF-SHA256)
|
|
KDF_ID = 0x0001 # HKDF-SHA256
|
|
AEAD_ID = 0x0003 # ChaCha20Poly1305
|
|
MODE_BASE = 0x00 # HPKE Base
|
|
SUITE_ID = b"HPKE" + KEM_ID.to_bytes(2, "big") + KDF_ID.to_bytes(2, "big") + AEAD_ID.to_bytes(2, "big")
|
|
N_ENC = 32 # X25519 public
|
|
N_PK = 32
|
|
N_SK = 32
|
|
N_NOISE = 12 # ChaCha nonce length in HPKE
|
|
N_K = 32 # key
|
|
AAD_DEFAULT = b"verae.npe.v1" # bind content to this system
|
|
|
|
|
|
def _hmac(key: bytes, data: bytes) -> bytes:
|
|
h = hmac.HMAC(key, hashes.SHA256()) # RFC HMAC
|
|
h.update(data)
|
|
return h.finalize()
|
|
|
|
|
|
def _extract(salt: bytes, ikm: bytes) -> bytes:
|
|
if not salt: # HMAC key
|
|
salt = b"\x00" * 32
|
|
return _hmac(salt, ikm)
|
|
|
|
|
|
def _expand(prk: bytes, info: bytes, length: int) -> bytes:
|
|
out = b"" # T(0)|T(1)...
|
|
t = b""
|
|
i = 1
|
|
while len(out) < length:
|
|
t = _hmac(prk, t + info + bytes([i]))
|
|
out += t
|
|
i += 1
|
|
return out[:length]
|
|
|
|
|
|
def _labeled_extract(salt: bytes, label: bytes, ikm: bytes) -> bytes:
|
|
labeled = b"HPKE-v1" + SUITE_ID + label + ikm # RFC 9180
|
|
return _extract(salt, labeled)
|
|
|
|
|
|
def _labeled_expand(prk: bytes, label: bytes, info: bytes, length: int) -> bytes:
|
|
labeled = length.to_bytes(2, "big") + b"HPKE-v1" + SUITE_ID + label + info
|
|
return _expand(prk, labeled, length)
|
|
|
|
|
|
def generate_keypair() -> Tuple[bytes, bytes]:
|
|
"""Return (raw_sk, raw_pk) 32+32 bytes."""
|
|
sk = X25519PrivateKey.generate() # CSPRNG
|
|
pk = sk.public_key()
|
|
return (
|
|
sk.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption()),
|
|
pk.public_bytes(Encoding.Raw, PublicFormat.Raw),
|
|
)
|
|
|
|
|
|
def _dh(sk: bytes, pk: bytes) -> bytes:
|
|
priv = X25519PrivateKey.from_private_bytes(sk)
|
|
pub = X25519PublicKey.from_public_bytes(pk)
|
|
return priv.exchange(pub)
|
|
|
|
|
|
def _extract_and_expand(dh: bytes, kem_context: bytes) -> bytes:
|
|
eae_prk = _labeled_extract(b"", b"eae_prk", dh)
|
|
return _labeled_expand(eae_prk, b"shared_secret", kem_context, 32)
|
|
|
|
|
|
def _key_schedule(shared: bytes, info: bytes) -> Tuple[bytes, bytes]:
|
|
psk_id_hash = _labeled_extract(b"", b"psk_id_hash", b"")
|
|
info_hash = _labeled_extract(b"", b"info_hash", info)
|
|
ks_ctx = bytes([MODE_BASE]) + psk_id_hash + info_hash
|
|
secret = _labeled_extract(shared, b"secret", b"")
|
|
key = _labeled_expand(secret, b"key", ks_ctx, N_K)
|
|
base_nonce = _labeled_expand(secret, b"base_nonce", ks_ctx, N_NOISE)
|
|
return key, base_nonce
|
|
|
|
|
|
def seal(recipient_pk: bytes, plaintext: bytes, aad: bytes = AAD_DEFAULT) -> Tuple[bytes, bytes]:
|
|
"""HPKE-Base seal. Returns (encapped_key, ciphertext)."""
|
|
eph_sk, eph_pk = generate_keypair() # ephemeral
|
|
dh = _dh(eph_sk, recipient_pk)
|
|
kem_context = eph_pk + recipient_pk
|
|
shared = _extract_and_expand(dh, kem_context)
|
|
key, nonce = _key_schedule(shared, b"")
|
|
ct = ChaCha20Poly1305(key).encrypt(nonce, plaintext, aad)
|
|
return eph_pk, ct
|
|
|
|
|
|
def open_ct(recipient_sk: bytes, enc: bytes, ciphertext: bytes, aad: bytes = AAD_DEFAULT) -> bytes:
|
|
"""HPKE-Base open. Raises on auth failure."""
|
|
rec_sk = X25519PrivateKey.from_private_bytes(recipient_sk)
|
|
rec_pk = rec_sk.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
|
|
dh = _dh(recipient_sk, enc)
|
|
kem_context = enc + rec_pk
|
|
shared = _extract_and_expand(dh, kem_context)
|
|
key, nonce = _key_schedule(shared, b"")
|
|
return ChaCha20Poly1305(key).decrypt(nonce, ciphertext, aad)
|