feat: Restricting a secret for a certain VO

Author: Robin-Van-de-Merghel

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

@@ -137,6 +137,11 @@ class BadPilotCredentialsError(GenericError):
     tail = ""
 
 
+class BadPilotVOError(GenericError):
+    head = "Bad VO"
+    tail = ""
+
+
 class SecretNotFoundError(GenericError):
     head = "Secret"
     tail = "not found"

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

@@ -8,6 +8,7 @@
 
 from diracx.core.exceptions import (
     BadPilotCredentialsError,
+    BadPilotVOError,
     CredentialsNotFoundError,
     OverusedSecretError,
     PilotAlreadyExistsError,
@@ -112,32 +113,40 @@ async def verify_pilot_secret(
             )
 
         # 5. Verify the secret counter
-        if secret["SecretGlobalUseCount"] + 1 > secret["SecretGlobalUseCountMax"]:
-            raise OverusedSecretError(
-                data={
-                    "pilot_stamp" "pilot_hashed_secret": pilot_hashed_secret,
-                    "secret_global_use_count": secret["SecretGlobalUseCount"],
-                    "secret_global_use_count_max": secret["SecretGlobalUseCountMax"],
-                }
-            )
-
-        # Now the pilot is authorized. Increment the counters (globally and locally).
+        # 5.1 Only check if the SecretGlobalUseCountMax is defined
+        # If not defined, there is an infinite use.
+        if secret["SecretGlobalUseCountMax"]:
+            # 5.2 Finite use, we check if we can still login
+            if secret["SecretGlobalUseCount"] + 1 > secret["SecretGlobalUseCountMax"]:
+                raise OverusedSecretError(
+                    data={
+                        "pilot_stamp" "pilot_hashed_secret": pilot_hashed_secret,
+                        "secret_global_use_count": secret["SecretGlobalUseCount"],
+                        "secret_global_use_count_max": secret[
+                            "SecretGlobalUseCountMax"
+                        ],
+                    }
+                )
+
+        # 6. Now the pilot is authorized, increment the counters (globally and locally).
         try:
-            # Increment the local count
+            # 6.1 Increment the local count
             await self.increment_pilot_local_secret_and_last_time_use(
                 pilot_secret_id=pilot_credentials["PilotSecretID"],
                 pilot_stamp=pilot_credentials["PilotStamp"],
             )
 
-            # Increment the global count
+            # 6.2 Increment the global count
             await self.increment_global_secret_use(
                 secret_id=pilot_credentials["PilotSecretID"]
             )
         except Exception as e:  # Generic, to catch it.
             # Should NOT happen
-            # Still caught in case of an error in the counters
+            # Wrapped in a try/catch to still catch in case of an error in the counters
             # Caught and raised here to avoid raising a 4XX error
-            raise DBInBadStateError(detail="This should not happen.") from e
+            raise DBInBadStateError(
+                detail="This should not happen. Pilot has credentials, but has a corrupted secret."
+            ) from e
 
     async def add_pilots_bulk(
         self,
@@ -175,6 +184,7 @@ async def add_pilots_credentials_bulk(
         self,
         pilot_stamps: list[str],
         pilot_hashed_secrets: list[str],
+        vo: str | None,
         pilot_secret_use_count_max: int = 1,
     ) -> list[dict]:
 
@@ -186,6 +196,7 @@ async def add_pilots_credentials_bulk(
         await self.insert_unique_secrets_bulk(
             hashed_secrets=pilot_hashed_secrets,
             secret_global_use_count_max=pilot_secret_use_count_max,
+            vo=vo,
         )
 
         # Get the secret ids to later associate them with pilots
@@ -207,13 +218,17 @@ async def add_pilots_credentials_bulk(
         return secrets
 
     async def insert_unique_secrets_bulk(
-        self, hashed_secrets: list[str], secret_global_use_count_max: int = 1
+        self,
+        hashed_secrets: list[str],
+        vo: str | None,
+        secret_global_use_count_max: int = 1,
     ):
         """Bulk insert secrets. Raises an error in case of a Integrity violation."""
         values = [
             {
                 "SecretGlobalUseCountMax": secret_global_use_count_max,
                 "HashedSecret": hashed_secret,
+                "SecretVO": vo,
             }
             for hashed_secret in hashed_secrets
         ]
@@ -234,13 +249,19 @@ async def insert_unique_secrets_bulk(
                 ) from e
 
             # Other errors to catch
-            raise DBInBadStateError("Engine Specific error not caught") from e
+            raise DBInBadStateError("Engine Specific error not caught" + str(e)) from e
 
     async def associate_pilots_with_secrets_bulk(
         self, pilot_to_secret_id_mapping_values: list[dict[str, Any]]
     ):
         """Bulk associate pilots with secrets. Raises an error in case of a Integrity violation."""
         # Better to give as a parameter pilot to secret associations, rather than associating here.
+
+        # First verify that pilots can access a certain secret
+        await self.verify_that_pilot_can_access_secret_bulk(
+            pilot_to_secret_id_mapping_values
+        )
+
         stmt = insert(PilotToSecretMapping).values(pilot_to_secret_id_mapping_values)
 
         try:
@@ -261,6 +282,52 @@ async def associate_pilots_with_secrets_bulk(
                     detail="at least one of these pilots already have a secret",
                 ) from e
 
+    async def verify_that_pilot_can_access_secret_bulk(
+        self, pilot_to_secret_id_mapping_values: list[dict[str, Any]]
+    ):
+        # 1. Extract unique pilot_stamps and secret_ids
+        pilot_stamps = [
+            entry["PilotStamp"] for entry in pilot_to_secret_id_mapping_values
+        ]
+        secret_ids = [
+            entry["PilotSecretID"] for entry in pilot_to_secret_id_mapping_values
+        ]
+
+        # 2. Bulk fetch pilot and secret info
+        pilots = await self.get_pilots_by_stamp_bulk(pilot_stamps)
+        secrets = await self.get_secrets_by_secret_ids_bulk(secret_ids)
+
+        # 3. Build lookup maps
+        pilot_vo_map = {pilot["PilotStamp"]: pilot["VO"] for pilot in pilots}
+        secret_vo_map = {secret["SecretID"]: secret["SecretVO"] for secret in secrets}
+
+        # 4. Validate access
+        bad_mapping = []
+
+        for mapping in pilot_to_secret_id_mapping_values:
+            pilot_stamp = mapping["PilotStamp"]
+            secret_id = mapping["PilotSecretID"]
+
+            pilot_vo = pilot_vo_map[pilot_stamp]
+            secret_vo = secret_vo_map[secret_id]
+
+            # If secret_vo is set to NULL, everybody can access it
+            if not secret_vo:
+                continue
+
+            # Access allowed only if VOs match or secret_vo is open (None)
+            if secret_vo is not None and pilot_vo != secret_vo:
+                bad_mapping.append(
+                    {
+                        "pilot_stamp": pilot_stamp,
+                        "given_vo": pilot_vo,
+                        "expected_vo": secret_vo,
+                    }
+                )
+
+        if bad_mapping:
+            raise BadPilotVOError(data={"bad_mapping": str(bad_mapping)})
+
     async def set_secret_expirations_bulk(
         self, secret_ids: list[int], pilot_secret_expiration_dates: list[DateTime]
     ):
@@ -288,6 +355,7 @@ async def get_pilots_by_stamp_bulk(self, pilot_stamps: list[str]) -> list[dict]:
             PilotAgents,
             PilotNotFoundError,
             "pilot_stamp",
+            "PilotStamp",
             pilot_stamps,
         )
 
@@ -300,6 +368,7 @@ async def get_pilots_credentials_by_stamps_bulk(
             PilotToSecretMapping,
             CredentialsNotFoundError,
             "pilot_stamp",
+            "PilotStamp",
             pilot_stamps,
         )
 
@@ -310,6 +379,18 @@ async def get_secrets_by_hashed_secrets_bulk(self, hashed_secrets: list[str]):
             PilotSecrets,
             SecretNotFoundError,
             "hashed_secret",
+            "HashedSecret",
             hashed_secrets,
             order_by=("secret_id", "asc"),
         )
+
+    async def get_secrets_by_secret_ids_bulk(self, secret_ids: list[int]):
+        """Bulk fetch secrets. Ensure all secrets are found, else raise an error."""
+        return await fetch_records_bulk_or_raises(
+            self.conn,
+            PilotSecrets,
+            SecretNotFoundError,
+            "secret_id",
+            "SecretID",
+            secret_ids,
+        )

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

@@ -76,6 +76,9 @@ class PilotSecrets(PilotAgentsDBBase):
     )
     secret_creation_time = DateNowColumn("SecretCreationDate")
     secret_expiration_date = NullColumn("SecretExpirationDate", DateTime(timezone=True))
+    # To authorize only pilots from a specific VO to access a secret
+    # Null VO => Can be used by everyone
+    secret_vo = NullColumn("SecretVO", String(128))
 
     __table_args__ = (UniqueConstraint("HashedSecret", name="uq_hashed_secret"),)
 

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

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import hashlib
-import re
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, Type
 
@@ -151,21 +150,12 @@ def rows_to_dicts(rows):
     return [dict(row._mapping) for row in rows]
 
 
-# For efficiency
-_SNAKE_SPLIT_RE = re.compile(r"_+")
-
-
-def snake_to_pascal(snake_str: str) -> str:
-    # Split the string using the precompiled regex
-    # Capitalize each word and join to form PascalCase
-    return "".join(word.capitalize() for word in _SNAKE_SPLIT_RE.split(snake_str))
-
-
 async def fetch_records_bulk_or_raises(
     conn: AsyncConnection,
     model: Any,  # Here, we currently must use `Any` because `declarative_base()` returns any
     missing_elements_error_cls: Type[GenericError],
-    column_to_use: str,
+    column_attribute_name: str,
+    column_name: str,
     elements_to_fetch: list,
     order_by: tuple[str, str] | None = None,
 ) -> list[dict]:
@@ -178,6 +168,7 @@ async def fetch_records_bulk_or_raises(
         self.conn,
         PilotAgents,
         PilotNotFound,
+        "pilot_id",
         "PilotID",
         [1,2,3]
     )
@@ -186,7 +177,7 @@ async def fetch_records_bulk_or_raises(
     assert elements_to_fetch
 
     # Get the column that needs to be in elements_to_fetch
-    column = getattr(model, column_to_use)
+    column = getattr(model, column_attribute_name)
 
     # Create the request
     stmt = select(model).with_for_update().where(column.in_(elements_to_fetch))
@@ -209,13 +200,12 @@ async def fetch_records_bulk_or_raises(
         raise DBInBadStateError(detail="Seems to have duplicates in the database.")
 
     # Checks if we have every elements we wanted
-    camel_case_column_to_use = snake_to_pascal(column_to_use)
-    found_keys = {row[camel_case_column_to_use] for row in results}
+    found_keys = {row[column_name] for row in results}
     missing = set(elements_to_fetch) - found_keys
 
     if missing:
         raise missing_elements_error_cls(
-            data={camel_case_column_to_use: str(missing)}, detail=str(missing)
+            data={column_name: str(missing)}, detail=str(missing)
         )
 
     return results