pop_crypto.py

"""
pop_crypto.py  --  primitive glue for the hardened PoP addon.

All primitives are standardized and vetted (Ed25519, X25519, HKDF-SHA256,
AES-256-GCM, SHA-256) via pyca/cryptography. Nothing here is rolled by hand.
"""

import os
import json
import base64
import hashlib
from urllib.parse import urlsplit, urlunsplit

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey, Ed25519PublicKey)
from cryptography.hazmat.primitives.asymmetric.x25519 import (
    X25519PrivateKey, X25519PublicKey)
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.exceptions import InvalidSignature

ALLOWED_SIG_ALGS = {"EdDSA"}          # L6: allowlist; never widen carelessly


# ---- encoding ----------------------------------------------------------------

def b64u(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def b64u_dec(s: str) -> bytes:
    return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))


def _j(obj) -> bytes:
    return json.dumps(obj, separators=(",", ":")).encode()


def sha256(b: bytes) -> bytes:
    return hashlib.sha256(b).digest()


# ---- JWK (Ed25519 OKP) + RFC 7638 thumbprint --------------------------------

def jwk_from_ed25519(pub: Ed25519PublicKey) -> dict:
    return {"kty": "OKP", "crv": "Ed25519", "x": b64u(pub.public_bytes_raw())}


def ed25519_from_jwk(jwk: dict) -> Ed25519PublicKey:
    if jwk.get("kty") != "OKP" or jwk.get("crv") != "Ed25519":
        raise ValueError("unsupported sig JWK")
    return Ed25519PublicKey.from_public_bytes(b64u_dec(jwk["x"]))


def jwk_thumbprint(jwk: dict) -> str:
    canonical = {"crv": jwk["crv"], "kty": jwk["kty"], "x": jwk["x"]}
    return b64u(sha256(json.dumps(canonical, separators=(",", ":"),
                                  sort_keys=True).encode()))


# ---- minimal JWS compact (EdDSA only) ---------------------------------------

def jws_sign(header: dict, payload: dict, sk: Ed25519PrivateKey) -> str:
    if header.get("alg") not in ALLOWED_SIG_ALGS:
        raise ValueError("alg not allowed")
    si = (b64u(_j(header)) + "." + b64u(_j(payload))).encode()
    return si.decode() + "." + b64u(sk.sign(si))


def jws_split(token: str):
    h_b64, p_b64, s_b64 = token.split(".")
    return (json.loads(b64u_dec(h_b64)), json.loads(b64u_dec(p_b64)),
            (h_b64 + "." + p_b64).encode(), b64u_dec(s_b64))


def jws_verify(token: str, pub: Ed25519PublicKey) -> dict:
    header, payload, si, sig = jws_split(token)
    if header.get("alg") not in ALLOWED_SIG_ALGS:     # L6: kills alg:none/confusion
        raise ValueError("alg not allowed")
    pub.verify(sig, si)                               # raises InvalidSignature
    return payload


# ---- H3: canonical HTU so client and server agree, computed from trusted req -

def canonical_htu(url: str) -> str:
    """Normalize: lowercase scheme/host, drop default port, drop query+fragment."""
    p = urlsplit(url)
    scheme = (p.scheme or "https").lower()
    host = (p.hostname or "").lower()
    default = {"http": 80, "https": 443}.get(scheme)
    netloc = host if (p.port in (None, default)) else f"{host}:{p.port}"
    path = p.path or "/"
    return urlunsplit((scheme, netloc, path, "", ""))


# ---- M1: channel binding value ----------------------------------------------

def channel_binding(channel_value: bytes) -> str:
    """Hash of a server-observed channel value (TLS exporter or client-cert DER)."""
    return b64u(sha256(channel_value))


# ---- M3: claim confidentiality, modeled on JWE ECDH-ES + A256GCM ------------
# (Faithful construction with vetted primitives; for wire interop use a JOSE lib.)

_ENC_AAD = b"jwt-pop-enc-v1"


def encrypt_to(recipient_x25519_pub: X25519PublicKey, plaintext: bytes) -> str:
    eph = X25519PrivateKey.generate()
    shared = eph.exchange(recipient_x25519_pub)
    key = HKDF(algorithm=hashes.SHA256(), length=32, salt=None,
               info=b"jwt-pop ecdh-es a256gcm").derive(shared)
    nonce = os.urandom(12)
    ct = AESGCM(key).encrypt(nonce, plaintext, _ENC_AAD)
    epk = {"kty": "OKP", "crv": "X25519",
           "x": b64u(eph.public_key().public_bytes_raw())}
    header = {"alg": "ECDH-ES", "enc": "A256GCM", "epk": epk}
    return "ENC1." + b64u(_j(header)) + "." + b64u(nonce) + "." + b64u(ct)


def decrypt_with(recipient_x25519_priv: X25519PrivateKey, token: str) -> bytes:
    tag, h_b64, n_b64, c_b64 = token.split(".")
    if tag != "ENC1":
        raise ValueError("not an encrypted token")
    header = json.loads(b64u_dec(h_b64))
    epk = X25519PublicKey.from_public_bytes(b64u_dec(header["epk"]["x"]))
    shared = recipient_x25519_priv.exchange(epk)
    key = HKDF(algorithm=hashes.SHA256(), length=32, salt=None,
               info=b"jwt-pop ecdh-es a256gcm").derive(shared)
    return AESGCM(key).decrypt(b64u_dec(n_b64), b64u_dec(c_b64), _ENC_AAD)


def is_encrypted(token: str) -> bool:
    return token.startswith("ENC1.")


def x25519_public_bytes(pub: X25519PublicKey) -> bytes:
    return pub.public_bytes_raw()