paskia/paskia/util/oidjwt.py

"""
OIDC JWT utilities for signing ID tokens and serving JWKS.
"""

import hashlib
import logging
from base64 import urlsafe_b64encode
from datetime import UTC, datetime, timedelta
from pathlib import Path
from uuid import UUID

import jwt
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

_logger = logging.getLogger(__name__)

# JWT signing key (loaded on first use)
_private_key = None
_public_key = None
_kid: str | None = None

# Key file location (same directory as database)
_KEY_FILE = Path("oidc_key.pem")


def _load_or_generate_key() -> None:
    """Load existing Ed25519 key or generate a new one."""
    global _private_key, _public_key, _kid

    if _KEY_FILE.exists():
        _logger.info("Loading OIDC signing key from %s", _KEY_FILE)
        pem_data = _KEY_FILE.read_bytes()
        _private_key = serialization.load_pem_private_key(pem_data, password=None)
    else:
        _logger.info("Generating new OIDC signing key")
        _private_key = Ed25519PrivateKey.generate()
        pem_data = _private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        _KEY_FILE.write_bytes(pem_data)
        _logger.info("Saved OIDC signing key to %s", _KEY_FILE)

    _public_key = _private_key.public_key()
    # Generate kid from public key fingerprint
    pub_der = _public_key.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    _kid = hashlib.sha256(pub_der).hexdigest()[:16]


def _ensure_key() -> None:
    """Ensure key is loaded."""
    if _private_key is None:
        _load_or_generate_key()


def get_jwks() -> dict:
    """Get JWKS (JSON Web Key Set) for public key verification."""
    _ensure_key()
    assert _public_key is not None
    # Ed25519 public key is 32 bytes raw
    pub_bytes = _public_key.public_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PublicFormat.Raw,
    )
    return {
        "keys": [
            {
                "kty": "OKP",
                "crv": "Ed25519",
                "use": "sig",
                "alg": "EdDSA",
                "kid": _kid,
                "x": urlsafe_b64encode(pub_bytes).rstrip(b"=").decode("ascii"),
            }
        ]
    }


def create_id_token(
    issuer: str,
    subject: UUID,
    audience: str,  # client_id
    nonce: str | None = None,
    sid: str | None = None,
    name: str | None = None,
    preferred_username: str | None = None,
    email: str | None = None,
    permissions: list[str] | None = None,
    auth_time: datetime | None = None,
    expires_in: int = 3600,
) -> str:
    """Create a signed ID token (JWT).

    Args:
        issuer: Token issuer (site URL)
        subject: User UUID (sub claim)
        audience: Client ID (aud claim)
        nonce: Nonce from authorization request
        sid: Session ID for backchannel logout
        name: User's display name
        preferred_username: User's preferred username
        email: User's email address
        permissions: List of permission scopes
        auth_time: When the user authenticated (last credential use time)
        expires_in: Token lifetime in seconds

    Returns:
        Signed JWT string
    """
    _ensure_key()
    now = datetime.now(UTC)
    payload = {
        "iss": issuer,
        "sub": str(subject),
        "aud": audience,
        "iat": int(now.timestamp()),
        "exp": int((now + timedelta(seconds=expires_in)).timestamp()),
    }
    if nonce:
        payload["nonce"] = nonce
    if sid:
        payload["sid"] = sid
    if name:
        payload["name"] = name
    if preferred_username:
        payload["preferred_username"] = preferred_username
    if email:
        payload["email"] = email
    if permissions:
        payload["permissions"] = permissions
    if auth_time:
        payload["auth_time"] = int(auth_time.timestamp())

    return jwt.encode(payload, _private_key, algorithm="EdDSA", headers={"kid": _kid})


def create_access_token(
    issuer: str,
    subject: UUID,
    audience: str,
    scope: str,
    expires_in: int = 3600,
) -> str:
    """Create a signed access token (JWT) for userinfo endpoint.

    Args:
        issuer: Token issuer (site URL)
        subject: User UUID
        audience: Client ID
        scope: Granted scopes
        expires_in: Token lifetime in seconds

    Returns:
        Signed JWT string
    """
    _ensure_key()
    now = datetime.now(UTC)
    payload = {
        "iss": issuer,
        "sub": str(subject),
        "aud": audience,
        "scope": scope,
        "iat": int(now.timestamp()),
        "exp": int((now + timedelta(seconds=expires_in)).timestamp()),
    }
    return jwt.encode(payload, _private_key, algorithm="EdDSA", headers={"kid": _kid})


def decode_access_token(
    token: str, issuer: str, audience: str | None = None
) -> dict | None:
    """Decode and verify an access token.

    Args:
        token: JWT string
        issuer: Expected issuer
        audience: Optional expected audience (client_id). If provided, aud claim must match.

    Returns:
        Decoded payload or None if invalid
    """
    _ensure_key()
    try:
        return jwt.decode(
            token,
            _public_key,
            algorithms=["EdDSA"],
            issuer=issuer,
            audience=audience,  # PyJWT handles None by skipping verification
        )
    except jwt.PyJWTError:
        return None