Merge pull request #499 from aldbr/main_FEAT_jwks-rotation

feat: introduce jwks

Author: chrisburr

diracx-client/src/diracx/client/_generated/aio/operations/_operations.py

@@ -54,6 +54,7 @@
     build_jobs_unassign_bulk_jobs_sandboxes_request,
     build_jobs_unassign_job_sandboxes_request,
     build_well_known_get_installation_metadata_request,
+    build_well_known_get_jwks_request,
     build_well_known_get_openid_configuration_request,
 )
 from .._configuration import DiracConfiguration
@@ -146,6 +147,57 @@ async def get_openid_configuration(
 
         return deserialized  # type: ignore
 
+    @distributed_trace_async
+    async def get_jwks(self, **kwargs: Any) -> Dict[str, Any]:
+        """Get Jwks.
+
+        Get the JWKs (public keys).
+
+        :return: dict mapping str to any
+        :rtype: dict[str, any]
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+        error_map: MutableMapping = {
+            401: ClientAuthenticationError,
+            404: ResourceNotFoundError,
+            409: ResourceExistsError,
+            304: ResourceNotModifiedError,
+        }
+        error_map.update(kwargs.pop("error_map", {}) or {})
+
+        _headers = kwargs.pop("headers", {}) or {}
+        _params = kwargs.pop("params", {}) or {}
+
+        cls: ClsType[Dict[str, Any]] = kwargs.pop("cls", None)
+
+        _request = build_well_known_get_jwks_request(
+            headers=_headers,
+            params=_params,
+        )
+        _request.url = self._client.format_url(_request.url)
+
+        _stream = False
+        pipeline_response: PipelineResponse = (
+            await self._client._pipeline.run(  # pylint: disable=protected-access
+                _request, stream=_stream, **kwargs
+            )
+        )
+
+        response = pipeline_response.http_response
+
+        if response.status_code not in [200]:
+            map_error(
+                status_code=response.status_code, response=response, error_map=error_map
+            )
+            raise HttpResponseError(response=response)
+
+        deserialized = self._deserialize("{object}", pipeline_response.http_response)
+
+        if cls:
+            return cls(pipeline_response, deserialized, {})  # type: ignore
+
+        return deserialized  # type: ignore
+
     @distributed_trace_async
     async def get_installation_metadata(self, **kwargs: Any) -> _models.Metadata:
         """Get Installation Metadata.

diracx-client/src/diracx/client/_generated/models/_models.py

@@ -563,6 +563,8 @@ class OpenIDConfiguration(_serialization.Model):
     :vartype device_authorization_endpoint: str
     :ivar revocation_endpoint: Revocation Endpoint. Required.
     :vartype revocation_endpoint: str
+    :ivar jwks_uri: Jwks Uri. Required.
+    :vartype jwks_uri: str
     :ivar grant_types_supported: Grant Types Supported. Required.
     :vartype grant_types_supported: list[str]
     :ivar scopes_supported: Scopes Supported. Required.
@@ -585,6 +587,7 @@ class OpenIDConfiguration(_serialization.Model):
         "authorization_endpoint": {"required": True},
         "device_authorization_endpoint": {"required": True},
         "revocation_endpoint": {"required": True},
+        "jwks_uri": {"required": True},
         "grant_types_supported": {"required": True},
         "scopes_supported": {"required": True},
         "response_types_supported": {"required": True},
@@ -603,6 +606,7 @@ class OpenIDConfiguration(_serialization.Model):
             "type": "str",
         },
         "revocation_endpoint": {"key": "revocation_endpoint", "type": "str"},
+        "jwks_uri": {"key": "jwks_uri", "type": "str"},
         "grant_types_supported": {"key": "grant_types_supported", "type": "[str]"},
         "scopes_supported": {"key": "scopes_supported", "type": "[str]"},
         "response_types_supported": {
@@ -632,6 +636,7 @@ def __init__(
         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],
@@ -653,6 +658,8 @@ def __init__(
         :paramtype device_authorization_endpoint: str
         :keyword revocation_endpoint: Revocation Endpoint. Required.
         :paramtype revocation_endpoint: str
+        :keyword jwks_uri: Jwks Uri. Required.
+        :paramtype jwks_uri: str
         :keyword grant_types_supported: Grant Types Supported. Required.
         :paramtype grant_types_supported: list[str]
         :keyword scopes_supported: Scopes Supported. Required.
@@ -675,6 +682,7 @@ def __init__(
         self.authorization_endpoint = authorization_endpoint
         self.device_authorization_endpoint = device_authorization_endpoint
         self.revocation_endpoint = revocation_endpoint
+        self.jwks_uri = jwks_uri
         self.grant_types_supported = grant_types_supported
         self.scopes_supported = scopes_supported
         self.response_types_supported = response_types_supported

diracx-client/src/diracx/client/_generated/operations/_operations.py

@@ -53,6 +53,20 @@ def build_well_known_get_openid_configuration_request(
     return HttpRequest(method="GET", url=_url, headers=_headers, **kwargs)
 
 
+def build_well_known_get_jwks_request(**kwargs: Any) -> HttpRequest:
+    _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
+
+    accept = _headers.pop("Accept", "application/json")
+
+    # Construct URL
+    _url = "/.well-known/jwks.json"
+
+    # Construct headers
+    _headers["Accept"] = _SERIALIZER.header("accept", accept, "str")
+
+    return HttpRequest(method="GET", url=_url, headers=_headers, **kwargs)
+
+
 def build_well_known_get_installation_metadata_request(
     **kwargs: Any,
 ) -> HttpRequest:  # pylint: disable=name-too-long
@@ -746,6 +760,57 @@ def get_openid_configuration(self, **kwargs: Any) -> _models.OpenIDConfiguration
 
         return deserialized  # type: ignore
 
+    @distributed_trace
+    def get_jwks(self, **kwargs: Any) -> Dict[str, Any]:
+        """Get Jwks.
+
+        Get the JWKs (public keys).
+
+        :return: dict mapping str to any
+        :rtype: dict[str, any]
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+        error_map: MutableMapping = {
+            401: ClientAuthenticationError,
+            404: ResourceNotFoundError,
+            409: ResourceExistsError,
+            304: ResourceNotModifiedError,
+        }
+        error_map.update(kwargs.pop("error_map", {}) or {})
+
+        _headers = kwargs.pop("headers", {}) or {}
+        _params = kwargs.pop("params", {}) or {}
+
+        cls: ClsType[Dict[str, Any]] = kwargs.pop("cls", None)
+
+        _request = build_well_known_get_jwks_request(
+            headers=_headers,
+            params=_params,
+        )
+        _request.url = self._client.format_url(_request.url)
+
+        _stream = False
+        pipeline_response: PipelineResponse = (
+            self._client._pipeline.run(  # pylint: disable=protected-access
+                _request, stream=_stream, **kwargs
+            )
+        )
+
+        response = pipeline_response.http_response
+
+        if response.status_code not in [200]:
+            map_error(
+                status_code=response.status_code, response=response, error_map=error_map
+            )
+            raise HttpResponseError(response=response)
+
+        deserialized = self._deserialize("{object}", pipeline_response.http_response)
+
+        if cls:
+            return cls(pipeline_response, deserialized, {})  # type: ignore
+
+        return deserialized  # type: ignore
+
     @distributed_trace
     def get_installation_metadata(self, **kwargs: Any) -> _models.Metadata:
         """Get Installation Metadata.

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,11 +19,12 @@
 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.jwk import KeySet, RSAKey
 from pydantic import (
+    AliasChoices,
     AnyUrl,
     BeforeValidator,
     Field,
@@ -51,28 +54,52 @@ 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())
-
-
-def _maybe_load_key_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()
+
+        # 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_keys_from_file(value: Any) -> Any:
+    """Load jwks from files if needed."""
+    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")
+            value = Path(url.path).read_text()
+
+    if isinstance(value, str) and value.strip().startswith("-----BEGIN"):
+        return json.dumps(
+            KeySet(
+                keys=[
+                    RSAKey.import_key(
+                        value,  # type: ignore
+                        parameters={"key_ops": ["sign", "verify"], "alg": "RS256"},  # type: ignore
+                    )
+                ]
+            ).as_dict(private=True)
+        )
     return value
 
 
-TokenSigningKey = Annotated[
-    _TokenSigningKey, BeforeValidator(_maybe_load_key_from_file)
+TokenSigningKeyStore = Annotated[
+    _TokenSigningKeyStore,
+    BeforeValidator(_maybe_load_keys_from_file),
 ]
 
 
@@ -124,7 +151,9 @@ def create(cls) -> Self:
 class AuthSettings(ServiceSettingsBase):
     """Settings for the authentication service."""
 
-    model_config = SettingsConfigDict(env_prefix="DIRACX_SERVICE_AUTH_")
+    model_config = SettingsConfigDict(
+        env_prefix="DIRACX_SERVICE_AUTH_", validate_by_name=True
+    )
 
     dirac_client_id: str = "myDIRACClientID"
     # TODO: This should be taken dynamically
@@ -137,8 +166,14 @@ class AuthSettings(ServiceSettingsBase):
     state_key: FernetKey
 
     token_issuer: str
-    token_key: TokenSigningKey
-    token_algorithm: str = "RS256"  # noqa: S105
+    token_keystore: TokenSigningKeyStore = Field(
+        validation_alias=AliasChoices(
+            "token_keystore",
+            "DIRACX_SERVICE_AUTH_TOKEN_KEYSTORE",
+            "DIRACX_SERVICE_AUTH_TOKEN_KEY",
+        )
+    )
+    token_allowed_algorithms: list[str] = ["RS256", "EdDSA"]  # noqa: S105
     access_token_expire_minutes: int = 20
     refresh_token_expire_minutes: int = 60