feat: A user can create secrets, and associate them to a pilot

Author: Robin-Van-de-Merghel

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

@@ -152,6 +152,11 @@ class CredentialsNotFoundError(GenericError):
     tail = "not found"
 
 
+class CredentialsAlreadyExistError(GenericError):
+    head = "Credentials"
+    tail = "already exist"
+
+
 class SecretHasExpiredError(GenericError):
     head = "Secret"
     tail = "has expired"

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

@@ -225,14 +225,17 @@ class TokenResponse(BaseModel):
     refresh_token: str | None = None
 
 
-class PilotCredentialsInfo(BaseModel):
-    pilot_stamp: str | None
+class PilotSecretsInfo(BaseModel):
     pilot_secret: str
     pilot_secret_expires_in: int
 
 
-class PilotCredentialsResponse(BaseModel):
-    pilot_credentials: list[PilotCredentialsInfo]
+class PilotStampInfo(BaseModel):
+    pilot_stamp: str
+
+
+class PilotCredentialsInfo(PilotSecretsInfo, PilotStampInfo):
+    pass
 
 
 class AccessTokenPayload(TokenPayload):

diracx-db/src/diracx/db/sql/pilot_agents/db.py

@@ -10,8 +10,8 @@
 from diracx.core.exceptions import (
     BadPilotCredentialsError,
     BadPilotVOError,
+    CredentialsAlreadyExistError,
     CredentialsNotFoundError,
-    PilotAlreadyExistsError,
     PilotNotFoundError,
     SecretAlreadyExistsError,
     SecretHasExpiredError,
@@ -87,25 +87,41 @@ async def verify_pilot_secret(
         pilots_credentials = await self.get_pilots_credentials_by_stamps_bulk(
             [pilot_stamp]
         )
-        pilot_credentials = pilots_credentials[
-            0
-        ]  # Semantic, assured by fetch_records_bulk_or_raises
 
         # 2. Get the pilot secret itself
         secrets = await self.get_secrets_by_hashed_secrets_bulk([pilot_hashed_secret])
         secret = secrets[0]  # Semantic, assured by fetch_records_bulk_or_raises
 
+        matches = [
+            pilot_credential
+            for pilot_credential in pilots_credentials
+            if secret["SecretID"] == pilot_credential["PilotSecretID"]
+        ]
+
         # 3. Compare the secret_id
-        if not secret["SecretID"] == pilot_credentials["PilotSecretID"]:
+        if len(matches) == 0:
 
             raise BadPilotCredentialsError(
                 data={
                     "pilot_stamp": pilot_stamp,
                     "pilot_hashed_secret": pilot_hashed_secret,
                     "real_hashed_secret": secret["HashedSecret"],
-                    str(secret["SecretID"]): str(pilot_credentials["PilotSecretID"]),
+                    "pilot_secret_id[]": str(
+                        [
+                            pilot_credential["PilotSecretID"]
+                            for pilot_credential in pilots_credentials
+                        ]
+                    ),
+                    "secret_id": secret["SecretID"],
+                    "test": str(pilots_credentials),
                 }
             )
+        elif len(matches) > 1:
+
+            raise DBInBadStateError(
+                detail="This should not happen. Duplicates in the database."
+            )
+        pilot_credentials = matches[0]  # Semantic
 
         # 4. Check if the secret is expired
         now = datetime.now(tz=timezone.utc)
@@ -260,20 +276,24 @@ async def associate_pilots_with_secrets_bulk(
         try:
             await self.conn.execute(stmt)
             await self.conn.commit()
+
         except (IntegrityError, OperationalError) as e:
             # Undo changes
             await self.conn.rollback()
-
             if "foreign key" in str(e.orig).lower():
                 raise PilotNotFoundError(
                     data={"pilot_stamps": str(pilot_to_secret_id_mapping_values)},
                     detail="at least one of these pilots or secrets does not exist",
                 ) from e
-            if "duplicate entry" in str(e.orig).lower():
-                raise PilotAlreadyExistsError(
+            if any(
+                el in str(e.orig).lower()
+                for el in ["duplicate entry", "unique constraint"]
+            ):
+                raise CredentialsAlreadyExistError(
                     data={"pilot_stamps": str(pilot_to_secret_id_mapping_values)},
                     detail="at least one of these pilots already have a secret",
                 ) from e
+            raise NotImplementedError(f"This error is not caught: {str(e.orig)}") from e
 
     async def verify_that_pilot_can_access_secret_bulk(
         self, pilot_to_secret_id_mapping_values: list[dict[str, Any]]
@@ -363,6 +383,7 @@ async def get_pilots_credentials_by_stamps_bulk(
             "pilot_stamp",
             "PilotStamp",
             pilot_stamps,
+            allow_more_than_one_result_per_input=True,
         )
 
     async def get_secrets_by_hashed_secrets_bulk(self, hashed_secrets: list[str]):

diracx-db/src/diracx/db/sql/utils/functions.py

@@ -158,6 +158,7 @@ async def fetch_records_bulk_or_raises(
     column_name: str,
     elements_to_fetch: list,
     order_by: tuple[str, str] | None = None,
+    allow_more_than_one_result_per_input: bool = False,
 ) -> list[dict]:
     """Fetches a list of elements in a table, returns a list of elements.
     All elements fro the `element_to_fetch` **must** be present.
@@ -196,8 +197,9 @@ async def fetch_records_bulk_or_raises(
     results = rows_to_dicts(await conn.execute(stmt))
 
     # Detects duplicates
-    if len(results) > len(elements_to_fetch):
-        raise DBInBadStateError(detail="Seems to have duplicates in the database.")
+    if not allow_more_than_one_result_per_input:
+        if len(results) > len(elements_to_fetch):
+            raise DBInBadStateError(detail="Seems to have duplicates in the database.")
 
     # Checks if we have every elements we wanted
     found_keys = {row[column_name] for row in results}

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

@@ -7,10 +7,11 @@
 from diracx.core.config import Config
 from diracx.core.exceptions import (
     ConfigurationError,
+    CredentialsAlreadyExistError,
     PilotAlreadyExistsError,
     PilotNotFoundError,
 )
-from diracx.core.models import PilotCredentialsInfo, PilotCredentialsResponse
+from diracx.core.models import PilotCredentialsInfo, PilotSecretsInfo, PilotStampInfo
 from diracx.core.settings import AuthSettings
 from diracx.db.sql import PilotAgentsDB
 
@@ -52,7 +53,7 @@ async def create_credentials(
         secret["SecretCreationDate"]
         + timedelta(
             seconds=(
-                expiration_minutes
+                expiration_minutes * 60
                 if expiration_minutes
                 else settings.pilot_secret_expire_seconds
             )
@@ -74,6 +75,36 @@ async def create_credentials(
     return random_secrets, hashed_secrets, expiration_dates_timestamps
 
 
+async def associate_pilots_with_secrets(
+    pilot_db: PilotAgentsDB,
+    pilot_stamps: list[str],
+    secrets: list[str] | None = None,
+    hashed_secrets: list[str] | None = None,
+):
+
+    if not hashed_secrets:
+        assert secrets
+        hashed_secrets = [hash(secret) for secret in secrets]
+
+    # Get the secret ids to later associate them with pilots
+    secrets_obj = await pilot_db.get_secrets_by_hashed_secrets_bulk(hashed_secrets)
+    secret_ids = [secret["SecretID"] for secret in secrets_obj]
+
+    if len(secret_ids) == 1:
+        secret_ids = secret_ids * len(pilot_stamps)
+
+    # Associates pilots with their secrets
+    pilot_to_secret_id_mapping_values = [
+        {
+            "PilotSecretID": secret_id,
+            "PilotStamp": pilot_stamp,
+        }
+        for pilot_stamp, secret_id in zip(pilot_stamps, secret_ids)
+    ]
+
+    await pilot_db.associate_pilots_with_secrets_bulk(pilot_to_secret_id_mapping_values)
+
+
 async def add_pilot_credentials(
     pilot_stamps: list[str],
     pilot_db: PilotAgentsDB,
@@ -91,29 +122,25 @@ async def add_pilot_credentials(
         )
     )
 
-    # Get the secret ids to later associate them with pilots
-    secrets = await pilot_db.get_secrets_by_hashed_secrets_bulk(hashed_secrets)
-    secret_ids = [secret["SecretID"] for secret in secrets]
-
-    # Associates pilots with their secrets
-    pilot_to_secret_id_mapping_values = [
-        {
-            "PilotSecretID": secret_id,
-            "PilotStamp": pilot_stamp,
-        }
-        for pilot_stamp, secret_id in zip(pilot_stamps, secret_ids)
-    ]
-    await pilot_db.associate_pilots_with_secrets_bulk(pilot_to_secret_id_mapping_values)
+    try:
+        await associate_pilots_with_secrets(
+            pilot_db=pilot_db, hashed_secrets=hashed_secrets, pilot_stamps=pilot_stamps
+        )
+    except CredentialsAlreadyExistError as e:
+        # Undo everything in case of an error.
+        # TODO: Validate in PR
+        await pilot_db.conn.rollback()
+        raise e
 
     return random_secrets, expiration_dates_timestamps
 
 
 def create_pilot_credentials_response(
-    pilot_stamps: list[str | None],
+    pilot_stamps: list[str],
     pilot_secrets: list[str],
     pilot_expiration_dates: list[int],
-) -> PilotCredentialsResponse:
-    credentials_list = [
+) -> list[PilotCredentialsInfo]:
+    return [
         PilotCredentialsInfo(
             pilot_stamp=pilot_stamp,
             pilot_secret=secret,
@@ -124,7 +151,22 @@ def create_pilot_credentials_response(
         )
     ]
 
-    return PilotCredentialsResponse(pilot_credentials=credentials_list)
+
+def create_secrets_response(
+    pilot_secrets: list[str],
+    pilot_expiration_dates: list[int],
+) -> list[PilotSecretsInfo]:
+    return [
+        PilotSecretsInfo(
+            pilot_secret=secret,
+            pilot_secret_expires_in=expires_in,
+        )
+        for secret, expires_in in zip(pilot_secrets, pilot_expiration_dates)
+    ]
+
+
+def create_stamp_response(pilot_stamps: list[str]) -> list[PilotStampInfo]:
+    return [PilotStampInfo(pilot_stamp=stamp) for stamp in pilot_stamps]
 
 
 def get_registry_and_group_configuration(config: Config, vo: str):
@@ -212,6 +254,8 @@ async def register_new_pilots(
                 pilots_that_already_exist = set(pilot_stamps) - set(
                     literal_eval(e.detail)
                 )
+            else:
+                raise ValueError("Bad internal error.")
         except AttributeError as e2:
             raise ValueError("Must be defined and a set string representation") from e2
 

diracx-routers/pyproject.toml

@@ -46,10 +46,12 @@ auth = "diracx.routers.auth:router"
 config = "diracx.routers.configuration:router"
 health = "diracx.routers.health:router"
 jobs = "diracx.routers.jobs:router"
+pilots = "diracx.routers.pilots:router"
 
 [project.entry-points."diracx.access_policies"]
 WMSAccessPolicy = "diracx.routers.jobs.access_policies:WMSAccessPolicy"
 SandboxAccessPolicy = "diracx.routers.jobs.access_policies:SandboxAccessPolicy"
+PilotCredentialsAccessPolicy = "diracx.routers.pilots.access_policies:PilotCredentialsAccessPolicy"
 
 # Minimum version of the client supported
 [project.entry-points."diracx.min_client_version"]

diracx-routers/src/diracx/routers/auth/__init__.py

@@ -5,7 +5,6 @@
 from .authorize_code_flow import router as authorize_code_flow_router
 from .device_flow import router as device_flow_router
 from .management import router as management_router
-from .pilot import router as pilot_router
 from .token import router as token_router
 from .utils import has_properties
 
@@ -14,6 +13,5 @@
 router.include_router(management_router)
 router.include_router(authorize_code_flow_router)
 router.include_router(token_router)
-router.include_router(pilot_router)
 
 __all__ = ["has_properties", "verify_dirac_access_token"]