feat: Add pilot registration (secret-exchange)

Author: Robin-Van-de-Merghel

diracx-client/src/diracx/client/patches/pilots/aio.py

@@ -16,6 +16,7 @@
 from azure.core.tracing.decorator_async import distributed_trace_async
 
 from ..._generated.aio.operations._operations import PilotsOperations as _PilotsOperations
+from ..._generated.models._models import PilotCredentialsInfo
 from .common import (
     make_search_body,
     make_summary_body,
@@ -43,7 +44,7 @@ async def summary(self, **kwargs: Unpack[SummaryKwargs]) -> list[dict[str, Any]]
         return await super().summary(**make_summary_body(**kwargs))
 
     @distributed_trace_async
-    async def add_pilot_stamps(self, **kwargs: Unpack[AddPilotStampsKwargs]) -> None:
+    async def add_pilot_stamps(self, **kwargs: Unpack[AddPilotStampsKwargs]) -> list[PilotCredentialsInfo] | None:
         """TODO"""
         return await super().add_pilot_stamps(**make_add_pilot_stamps_body(**kwargs))
 

diracx-client/src/diracx/client/patches/pilots/common.py

@@ -99,6 +99,8 @@ class AddPilotStampsBody(TypedDict, total=False):
     pilot_references: dict[str, str]
     pilot_status: PilotStatus
     vo: str
+    generate_secrets: bool
+    pilot_secret_use_count_max: int | None
 
 class AddPilotStampsKwargs(AddPilotStampsBody, ResponseExtra): ...
 
@@ -112,7 +114,7 @@ def make_add_pilot_stamps_body(**kwargs: Unpack[AddPilotStampsKwargs]) -> Underl
     for key in AddPilotStampsBody.__optional_keys__:
         if key not in kwargs:
             continue
-        key = cast(Literal["pilot_stamps", "grid_type", "grid_site", "pilot_references", "pilot_status", "vo"], key)
+        key = cast(Literal["pilot_stamps", "grid_type", "grid_site", "pilot_references", "pilot_status", "vo", "generate_secrets", "pilot_secret_use_count_max"], key)
         value = kwargs.pop(key)
         if value is not None:
             body[key] = value

diracx-client/src/diracx/client/patches/pilots/sync.py

@@ -16,6 +16,7 @@
 from azure.core.tracing.decorator import distributed_trace
 
 from ..._generated.operations._operations import PilotsOperations as _PilotsOperations
+from ..._generated.models._models import PilotCredentialsInfo
 from .common import (
     make_search_body,
     make_summary_body,
@@ -43,7 +44,7 @@ def summary(self, **kwargs: Unpack[SummaryKwargs]) -> list[dict[str, Any]]:
         return super().summary(**make_summary_body(**kwargs))
 
     @distributed_trace
-    def add_pilot_stamps(self, **kwargs: Unpack[AddPilotStampsKwargs]) -> None:
+    def add_pilot_stamps(self, **kwargs: Unpack[AddPilotStampsKwargs]) -> list[PilotCredentialsInfo] | None:
         """TODO"""
         return super().add_pilot_stamps(**make_add_pilot_stamps_body(**kwargs))
 

diracx-core/src/diracx/core/exceptions.py

@@ -99,6 +99,9 @@ def __init__(self, job_id, detail: str = ""):
         )
 
 
+class BadTokenError(DiracError): ...
+
+
 class NotReadyError(DiracError):
     """Tried to access a value which is asynchronously loaded but not yet available."""
 
@@ -113,3 +116,15 @@ class PilotAlreadyExistsError(DiracError):
 
 class PilotAlreadyAssociatedWithJobError(DiracError):
     """We can't associate a pilot with the same job twice."""
+
+
+class SecretHasExpiredError(DiracError):
+    """If a secret expired."""
+
+
+class SecretNotFoundError(DiracError):
+    """If a secret not found."""
+
+
+class BadPilotCredentialsError(DiracError):
+    """If a pilot tries to auth with another pilot's credentials."""

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

@@ -5,12 +5,15 @@
 
 from __future__ import annotations
 
+import uuid as std_uuid
 from datetime import datetime
-from enum import StrEnum
-from typing import Literal, Optional
+from enum import StrEnum, auto
+from typing import Any, Literal, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, GetCoreSchemaHandler, GetJsonSchemaHandler
+from pydantic_core import CoreSchema, core_schema
 from typing_extensions import TypedDict
+from uuid_utils import UUID as _UUID
 
 
 class ScalarSearchOperator(StrEnum):
@@ -37,7 +40,7 @@ class ScalarSearchSpec(TypedDict):
 class VectorSearchSpec(TypedDict):
     parameter: str
     operator: VectorSearchOperator
-    values: list[str] | list[int]
+    values: list[str] | list[int] | list[bytes]
 
 
 SearchSpec = ScalarSearchSpec | VectorSearchSpec
@@ -179,6 +182,29 @@ class SandboxUploadResponse(BaseModel):
     fields: dict[str, str] = {}
 
 
+class UUID(_UUID):
+    """Subclass of uuid_utils.UUID to add pydantic support."""
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        """Use the stdlib uuid.UUID schema for validation and serialization."""
+        std_schema = handler(std_uuid.UUID)
+
+        def to_uuid_utils(u: std_uuid.UUID) -> UUID:
+            return cls(str(u))
+
+        return core_schema.no_info_after_validator_function(to_uuid_utils, std_schema)
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler
+    ) -> dict[str, Any]:
+        """Return the stdlib uuid.UUID schema for JSON serialization."""
+        return handler(core_schema)
+
+
 class GrantType(StrEnum):
     """Grant types for OAuth2."""
 
@@ -213,9 +239,14 @@ class OpenIDConfiguration(TypedDict):
     code_challenge_methods_supported: list[str]
 
 
-class TokenPayload(TypedDict):
+class BaseTokenPayload(TypedDict):
+    """This class helps having pilot and user tokens without code duplication."""
+
     jti: str
     exp: datetime
+
+
+class TokenPayload(BaseTokenPayload):
     dirac_policies: dict
 
 
@@ -308,3 +339,56 @@ class PilotStatus(StrEnum):
     ABORTED = "Aborted"
     #: Cannot get information about the pilot status:
     UNKNOWN = "Unknown"
+
+
+class PilotSecretConstraints(TypedDict, total=False):
+    VOs: list[str]  # Authorize only a list of VOs
+    PilotStamps: list[str]  # Authorize only a list of stamps
+    Sites: list[str]  # Authorize only a list of sites
+    # ...
+    # We can add constraints here
+
+
+class TokenType(StrEnum):
+    # Pilot token
+    PILOT_TOKEN = auto()
+    # User token
+    USER_TOKEN = auto()
+
+
+class PilotSecretsInfo(BaseModel):
+    pilot_secret: str
+    pilot_secret_expires_in: int
+
+
+class PilotAccessTokenPayload(BaseTokenPayload):
+    sub: str
+    vo: str
+    iss: str
+    pilot_stamp: str
+
+
+class PilotInfo(BaseModel):
+    pilot_stamp: str
+    vo: str
+    sub: str
+
+
+class PilotRefreshTokenPayload(BaseTokenPayload):
+    legacy_exchange: bool
+
+
+class PilotCredentialsInfo(PilotSecretsInfo):
+    pilot_stamp: str
+
+
+class PilotAuthCredentials(TypedDict):
+    pilot_stamp: str
+    pilot_secret: str
+
+
+class VacuumPilotAuth(PilotAuthCredentials):
+    vo: str
+    grid_type: str
+    grid_site: str
+    status: str

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

@@ -163,6 +163,8 @@ class AuthSettings(ServiceSettingsBase):
     token_allowed_algorithms: list[str] = ["RS256", "EdDSA"]  # noqa: S105
     access_token_expire_minutes: int = 20
     refresh_token_expire_minutes: int = 60
+    pilot_secret_expire_seconds: int = 3600
+    pilot_refresh_token_expire_hours: int = 168
 
     available_properties: set[SecurityProperty] = Field(
         default_factory=SecurityProperty.available_properties

diracx-core/src/diracx/core/utils.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from uuid import UUID
+
 __all__ = [
     "dotenv_files_from_environment",
     "serialize_credentials",
@@ -19,7 +21,7 @@
 from concurrent.futures import Future, ThreadPoolExecutor, wait
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
-from typing import Any, AsyncIterable, TypeVar
+from typing import Any, AsyncIterable, Mapping, TypeVar, cast
 
 from cachetools import Cache, TTLCache
 
@@ -271,3 +273,31 @@ async def batched_async(
         if strict and len(batch) != n:
             raise ValueError("batched(): incomplete batch")
         yield tuple(batch)
+
+
+def extract_timestamp_from_uuid7(uuid_str: str) -> datetime:
+    u = UUID(uuid_str)
+    ts_bytes = u.bytes[0:6]  # First 48 bits = timestamp in ms
+    timestamp_ms = int.from_bytes(ts_bytes, byteorder="big")
+    # Convert into seconds then to datetime
+    return datetime.fromtimestamp(timestamp_ms / 1000, timezone.utc)
+
+
+T_DICTS = TypeVar("T_DICTS", bound=Mapping[str, Any])
+
+
+def recursive_dict_merge(x: T_DICTS, y: T_DICTS) -> T_DICTS:
+    result: dict[str, Any] = dict(x)
+
+    for k, v in y.items():
+        if k in result:
+            if isinstance(result[k], dict) and isinstance(v, dict):
+                result[k] = recursive_dict_merge(result[k], v)
+            elif isinstance(result[k], list) and isinstance(v, list):
+                result[k] = result[k] + v
+            else:
+                result[k] = v
+        else:
+            result[k] = v
+
+    return cast(T_DICTS, result)