feat: introduce jwks

Author: aldbr

diracx-core/pyproject.toml

@@ -14,11 +14,11 @@ classifiers = [
 ]
 dependencies = [
     "aiobotocore>=2.15",
-    "authlib",
     "botocore>=1.35",
     "cachetools",
     "email_validator",
     "gitpython",
+    "joserfc",
     "pydantic >=2.10",
     "pydantic-settings",
     "pyyaml",

diracx-core/src/diracx/core/models.py

@@ -202,6 +202,7 @@ class OpenIDConfiguration(TypedDict):
     authorization_endpoint: str
     device_authorization_endpoint: str
     revocation_endpoint: str
+    jwks_uri: str
     grant_types_supported: list[str]
     scopes_supported: list[str]
     response_types_supported: list[str]

diracx-core/src/diracx/core/settings.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import json
+
 from diracx.core.properties import SecurityProperty
 from diracx.core.s3 import s3_bucket_exists
 
@@ -17,10 +19,10 @@
 from typing import TYPE_CHECKING, Annotated, Any, Self, TypeVar
 
 from aiobotocore.session import get_session
-from authlib.jose import JsonWebKey
 from botocore.config import Config
 from botocore.errorfactory import ClientError
 from cryptography.fernet import Fernet
+from joserfc.jws import KeySet
 from pydantic import (
     AnyUrl,
     BeforeValidator,
@@ -51,28 +53,40 @@ class SqlalchemyDsn(AnyUrl):
     )
 
 
-class _TokenSigningKey(SecretStr):
-    jwk: JsonWebKey
+class _TokenSigningKeyStore(SecretStr):
+    jwks: KeySet
 
     def __init__(self, data: str):
         super().__init__(data)
-        self.jwk = JsonWebKey.import_key(self.get_secret_value())
+
+        # Load the keys from the JSON string
+        try:
+            keys = json.loads(self.get_secret_value())
+        except json.JSONDecodeError as e:
+            raise ValueError("Invalid JSON string") from e
+        if not isinstance(keys, dict):
+            raise ValueError("Invalid JSON string")
+        self.jwks = KeySet.import_key_set(keys)  # type: ignore
 
 
-def _maybe_load_key_from_file(value: Any) -> Any:
+def _maybe_load_keys_from_file(value: Any) -> Any:
     """Load private keys from files if needed."""
-    if isinstance(value, str) and not value.strip().startswith("-----BEGIN"):
-        url = TypeAdapter(LocalFileUrl).validate_python(value)
-        if not url.scheme == "file":
-            raise ValueError("Only file:// URLs are supported")
-        if url.path is None:
-            raise ValueError("No path specified")
-        value = Path(url.path).read_text()
+    if isinstance(value, str):
+        # If the value is a string, we need to check if it is a JSON string or a file URL
+        if not (value.strip().startswith("{") or value.startswith("[")):
+            # If it is not a JSON string, we assume it is a file URL
+            url = TypeAdapter(LocalFileUrl).validate_python(value)
+            if not url.scheme == "file":
+                raise ValueError("Only file:// URLs are supported")
+            if url.path is None:
+                raise ValueError("No path specified")
+            return Path(url.path).read_text()
+
     return value
 
 
-TokenSigningKey = Annotated[
-    _TokenSigningKey, BeforeValidator(_maybe_load_key_from_file)
+TokenSigningKeyStore = Annotated[
+    _TokenSigningKeyStore, BeforeValidator(_maybe_load_keys_from_file)
 ]
 
 
@@ -137,8 +151,8 @@ class AuthSettings(ServiceSettingsBase):
     state_key: FernetKey
 
     token_issuer: str
-    token_key: TokenSigningKey
-    token_algorithm: str = "RS256"  # noqa: S105
+    token_keystore: TokenSigningKeyStore
+    token_allowed_algorithms: list[str] = ["RS256", "EdDSA"]  # noqa: S105
     access_token_expire_minutes: int = 20
     refresh_token_expire_minutes: int = 60
 

diracx-core/tests/test_secrets.py

@@ -1,47 +1,45 @@
 from __future__ import annotations
 
-from cryptography.hazmat.primitives import serialization
-from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
+import json
+
+from joserfc.jwk import KeySet, OKPKey
 from pydantic import TypeAdapter
+from uuid_utils import uuid7
 
-from diracx.core.settings import TokenSigningKey
+from diracx.core.settings import TokenSigningKeyStore
 
 
-def compare_keys(key1, key2):
-    """Compare two keys by checking their public keys."""
-    key1_public = key1.public_key().public_bytes(
-        encoding=serialization.Encoding.PEM,
-        format=serialization.PublicFormat.SubjectPublicKeyInfo,
-    )
-    key2_public = key2.public_key().public_bytes(
-        encoding=serialization.Encoding.PEM,
-        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+def test_token_signing_key(tmp_path):
+    keyset = KeySet(
+        keys=[
+            OKPKey.generate_key(
+                parameters={
+                    "key_ops": ["sign", "verify"],
+                    "alg": "EdDSA",
+                    "kid": uuid7().hex,
+                }
+            )
+        ]
     )
-    assert key1_public == key2_public
 
+    jwks_file = tmp_path / "jwks.json"
+    jwks_file.write_text(json.dumps(keyset.as_dict(private=True)))
 
-def test_token_signing_key(tmp_path):
-    private_key = Ed25519PrivateKey.generate()
-    private_key_pem = private_key.private_bytes(
-        encoding=serialization.Encoding.PEM,
-        format=serialization.PrivateFormat.PKCS8,
-        encryption_algorithm=serialization.NoEncryption(),
-    ).decode("ascii")
-    key_file = tmp_path / "private_key.pem"
-    key_file.write_text(private_key_pem)
-
-    adapter = TypeAdapter(TokenSigningKey)
-
-    # Test that we can load a key from a file
-    compare_keys(
-        adapter.validate_python(f"{key_file}").jwk.get_private_key(), private_key
+    adapter = TypeAdapter(TokenSigningKeyStore)
+
+    # Test that we can load a keystore from a file
+    assert (
+        adapter.validate_python(f"{jwks_file}").jwks.keys[0].kid == keyset.keys[0].kid
     )
-    compare_keys(
-        adapter.validate_python(f"file://{key_file}").jwk.get_private_key(),
-        private_key,
+    assert (
+        adapter.validate_python(f"file://{jwks_file}").jwks.keys[0].kid
+        == keyset.keys[0].kid
     )
 
-    # Test with can load the PEM data directly
-    compare_keys(
-        adapter.validate_python(private_key_pem).jwk.get_private_key(), private_key
+    # Test with can load the keystore data directly from a JSON string
+    assert (
+        adapter.validate_python(json.dumps(keyset.as_dict(private=True)))
+        .jwks.keys[0]
+        .kid
+        == keyset.keys[0].kid
     )

diracx-logic/pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "dirac",
     "diracx-core",
     "diracx-db",
+    "joserfc",
     "pydantic >=2.10",
     "uuid-utils",
 ]

diracx-logic/src/diracx/logic/__main__.py

@@ -0,0 +1,126 @@
+"""JWKS key management scripts.
+
+See https://datatracker.ietf.org/doc/html/rfc7517 for further details.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import logging
+from pathlib import Path
+
+from joserfc.jwk import JWKRegistry, Key, KeySet
+from uuid_utils import uuid7
+
+logger = logging.getLogger(__name__)
+
+# ---------- Helpers ----------------------------------------------------------
+
+
+def load_jwks(path: Path) -> KeySet:
+    """Return a (possibly empty) JWKSet."""
+    if path.exists():
+        return KeySet.import_key_set(json.loads(path.read_text()))
+    logger.warning("JWKS file %s not found – creating a new one", path)
+    return KeySet(keys=[])
+
+
+def save_jwks(path: Path, jwks: KeySet) -> None:
+    """Write JWKSet to disk *including* private parts."""
+    path.write_text(json.dumps(jwks.as_dict(private=True), indent=2))
+    logger.info("JWKS written to %s", path)
+
+
+def new_key(
+    kty: str = "OKP",
+    crv_or_size: str | int = "Ed25519",
+) -> Key:
+    """Create a fresh private signing key."""
+    parameters = {
+        "key_ops": ["sign", "verify"],
+        "alg": "EdDSA",
+        "kid": uuid7().hex,
+    }
+    return JWKRegistry.generate_key(
+        key_type=kty, crv_or_size=crv_or_size, private=True, parameters=parameters  # type: ignore[arg-type]
+    )
+
+
+# ---------- CLI --------------------------------------------------------------
+
+
+async def rotate_jwk(args):
+    """Rotate keys in a JWKS file by inserting a new key at index 0 (active)."""
+    logger.info("Rotating JWKs...")
+
+    crv_or_size = args.crv_or_size
+    if isinstance(crv_or_size, str) and crv_or_size.isdigit():
+        crv_or_size = int(crv_or_size)
+
+    jwks_path = Path(args.jwks_path)
+    jwks = load_jwks(jwks_path)
+
+    # Current key (at index 0) is set to "verify" only
+    if len(jwks.keys) > 0:
+        active_key = jwks.keys[0]
+        active_key_dict = active_key.as_dict(private=True)
+        active_key_dict["key_ops"] = sorted(
+            set(active_key_dict.get("key_ops", [])) - {"sign"}
+        )
+        jwks.keys[0] = JWKRegistry.import_key(active_key_dict)
+
+    jwk = new_key(args.kty, crv_or_size)
+    jwks.keys.insert(0, jwk)
+
+    save_jwks(jwks_path, jwks)
+
+
+async def delete_jwk(args):
+    """Delete a JWK from a JWKS file."""
+    logger.info("Deleting JWK...")
+
+    path = Path(args.jwks_path)
+    jwks = load_jwks(path)
+    jwks.keys = [k for k in jwks.keys if k.get("kid") != args.kid]
+    save_jwks(path, jwks)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    rotate_jwk_parser = subparsers.add_parser(
+        "rotate-jwk", help="Rotate JWK keys in a JWKS file"
+    )
+    rotate_jwk_parser.add_argument(
+        "--jwks-path", required=True, help="Path to the existing (old) JWKS JSON file."
+    )
+
+    rotate_jwk_parser.add_argument(
+        "--kty", default="OKP", help="Key type for the new key."
+    )
+    rotate_jwk_parser.add_argument(
+        "--crv-or-size", default="Ed25519", help="Curve or size for the new key."
+    )
+    rotate_jwk_parser.set_defaults(func=rotate_jwk)
+
+    delete_jwk_parser = subparsers.add_parser(
+        "delete-jwk", help="Delete a JWK key from a JWKS file"
+    )
+    delete_jwk_parser.add_argument(
+        "--jwks-path", required=True, help="Path to the JWKS JSON file."
+    )
+    delete_jwk_parser.add_argument(
+        "--kid", required=True, help="Key ID (kid) of the key to delete."
+    )
+    delete_jwk_parser.set_defaults(func=delete_jwk)
+
+    args = parser.parse_args()
+    logger.setLevel(logging.INFO)
+    asyncio.run(args.func(args))
+
+
+if __name__ == "__main__":
+    parse_args()

diracx-logic/src/diracx/logic/auth/token.py

@@ -6,8 +6,10 @@
 import hashlib
 import re
 from datetime import datetime, timedelta, timezone
+from typing import cast
 
-from authlib.jose import JsonWebToken
+from joserfc import jwt
+from joserfc.jwt import Claims
 from uuid_utils import UUID, uuid7
 
 from diracx.core.config import Config
@@ -356,11 +358,24 @@ async def exchange_token(
 
 
 def create_token(payload: TokenPayload, settings: AuthSettings) -> str:
-    jwt = JsonWebToken(settings.token_algorithm)
-    encoded_jwt = jwt.encode(
-        {"alg": settings.token_algorithm}, payload, settings.token_key.jwk
+    """Create a JWT token with the given payload and settings."""
+    signing_key = None
+    for key in settings.token_keystore.jwks.keys:
+        # TODO: https://github.com/authlib/joserfc/issues/52
+        key_ops = cast(list[str] | None, key.get("key_ops"))
+        if key_ops and "sign" in key_ops:
+            signing_key = key
+            break
+
+    if not signing_key:
+        raise ValueError("No signing key found in JWKS")
+
+    return jwt.encode(
+        header={"alg": signing_key.get("alg"), "kid": signing_key.get("kid")},
+        claims=cast(Claims, payload),
+        key=settings.token_keystore.jwks,
+        algorithms=settings.token_allowed_algorithms,
     )
-    return encoded_jwt.decode("ascii")
 
 
 async def insert_refresh_token(