fix: Moved secrets table to the AuthDB

Author: Robin-Van-de-Merghel

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

@@ -1,21 +1,26 @@
 from __future__ import annotations
 
 import secrets
+from typing import Any
 
-from sqlalchemy import insert, select, update
+from sqlalchemy import DateTime, bindparam, delete, insert, select, update
 from sqlalchemy.exc import IntegrityError, NoResultFound
 from uuid_utils import UUID, uuid7
 
 from diracx.core.exceptions import (
     AuthorizationError,
+    SecretNotFoundError,
     TokenNotFoundError,
 )
+from diracx.core.models import PilotSecretConstraints, SearchSpec, SortSpec
 from diracx.db.sql.utils import BaseSQLDB, hash, substract_date
+from diracx.db.sql.utils.functions import utcnow
 
 from .schema import (
     AuthorizationFlows,
     DeviceFlows,
     FlowStatus,
+    PilotSecrets,
     RefreshTokens,
     RefreshTokenStatus,
 )
@@ -264,3 +269,153 @@ async def revoke_user_refresh_tokens(self, subject):
             .where(RefreshTokens.sub == subject)
             .values(status=RefreshTokenStatus.REVOKED)
         )
+
+    # ------------- Pilot secrets mechanism -------------
+
+    async def insert_unique_secrets(
+        self,
+        hashed_secrets: list[bytes],
+        secret_global_use_count_max: int | None = 1,
+        secret_constraints: dict[bytes, PilotSecretConstraints] = {},
+    ):
+        """Bulk insert secrets.
+
+        Raises:
+        - NotImplementedError if we have an IntegrityError not caught
+
+        """
+        values = [
+            {
+                "SecretUUID": str(uuid7()),
+                "SecretRemainingUseCount": secret_global_use_count_max,
+                "HashedSecret": hashed_secret,
+                "SecretConstraints": secret_constraints.get(hashed_secret, {}),
+            }
+            for hashed_secret in hashed_secrets
+        ]
+
+        stmt = insert(PilotSecrets).values(values)
+        await self.conn.execute(stmt)
+
+    async def delete_secrets(self, secret_uuids: list[str]):
+        """Bulk delete secrets.
+
+        Raises SecretNotFoundError if one of the secret was not found.
+        """
+        stmt = delete(PilotSecrets).where(PilotSecrets.secret_uuid.in_(secret_uuids))
+
+        res = await self.conn.execute(stmt)
+
+        if res.rowcount != len(secret_uuids):
+            raise SecretNotFoundError(
+                "At least one of the secret has not been deleted."
+            )
+
+        # We NEED to commit here, because we will raise an error after this function
+        await self.conn.commit()
+
+    async def update_pilot_secret_use_time(self, secret_uuid: str) -> None:
+        """Updates when a pilot uses a secret.
+
+        Raises PilotNotFoundError if the pilot does not exist
+
+        """
+        #  Prepare the update statement
+        stmt = (
+            update(PilotSecrets)
+            .values(
+                pilot_secret_use_date=utcnow(),
+                secret_remaining_use_count=PilotSecrets.secret_remaining_use_count - 1,
+            )
+            .where(PilotSecrets.secret_uuid == secret_uuid)
+        )
+
+        # Execute the update using the connection
+        res = await self.conn.execute(stmt)
+
+        if res.rowcount == 0:
+            raise SecretNotFoundError("Unknown secret")
+
+    async def update_pilot_secrets_constraints(
+        self, hashed_secrets_to_pilot_stamps_mapping: list[dict[str, Any]]
+    ):
+        """Bulk associate pilots with secrets by updating theirs constraints.
+
+        Important: We have to provide the updated constraints.
+
+        Raises:
+        - PilotNotFoundError if one of the pilot does not exist
+        - NotImplementedError if at least of the pilot
+
+        """
+        # Better to give as a parameter pilot to secret associations, rather than associating here.
+
+        stmt = (
+            update(PilotSecrets)
+            .where(PilotSecrets.hashed_secret == bindparam("PilotHashedSecret"))
+            .values({"SecretConstraints": bindparam("PilotSecretConstraints")})
+        )
+
+        try:
+            await self.conn.execute(stmt, hashed_secrets_to_pilot_stamps_mapping)
+        except IntegrityError as e:
+            if "foreign key" in str(e.orig).lower():
+                raise SecretNotFoundError(
+                    detail="at least one of these secrets does not exist",
+                ) from e
+            raise NotImplementedError(f"This error is not caught: {str(e.orig)}") from e
+
+    async def set_secret_expirations(
+        self, secret_uuids: list[str], pilot_secret_expiration_dates: list[DateTime]
+    ):
+        """Bulk set expiration dates to secrets.
+
+        Raises:
+        - SecretNotFoundError if one of the secret_uuid is not associated with a secret.
+        - NotImplementedError if a integrity error is not caught.
+        -
+
+        """
+        values = [
+            {"b_SecretUUID": secret_uuid, "SecretExpirationDate": pilot_secret}
+            for secret_uuid, pilot_secret in zip(
+                secret_uuids, pilot_secret_expiration_dates
+            )
+        ]
+
+        #  Prepare the update statement
+        stmt = (
+            update(PilotSecrets)
+            .where(PilotSecrets.secret_uuid == bindparam("b_SecretUUID"))
+            .values({"SecretExpirationDate": bindparam("SecretExpirationDate")})
+        )
+
+        try:
+            await self.conn.execute(stmt, values)
+        except IntegrityError as e:
+            if "foreign key" in str(e.orig).lower():
+                raise SecretNotFoundError(
+                    detail="at least one of these secrets does not exist",
+                ) from e
+            raise NotImplementedError(f"This error is not caught: {str(e.orig)}") from e
+
+    async def search_secrets(
+        self,
+        parameters: list[str] | None,
+        search: list[SearchSpec],
+        sorts: list[SortSpec],
+        *,
+        distinct: bool = False,
+        per_page: int = 100,
+        page: int | None = None,
+    ) -> tuple[int, list[dict[Any, Any]]]:
+        """Search for secrets in the database."""
+        return await self._search(
+            table=PilotSecrets,
+            parameters=parameters,
+            search=search,
+            sorts=sorts,
+            distinct=distinct,
+            per_page=per_page,
+            page=page,
+        )

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

@@ -3,9 +3,13 @@
 from enum import Enum, auto
 
 from sqlalchemy import (
+    BINARY,
     JSON,
+    DateTime,
     Index,
+    SmallInteger,
     String,
+    UniqueConstraint,
     Uuid,
 )
 from sqlalchemy.orm import declarative_base
@@ -99,3 +103,28 @@ class RefreshTokens(Base):
     sub = Column("Sub", String(256), index=True)
 
     __table_args__ = (Index("index_status_sub", status, sub),)
+
+
+class PilotSecrets(Base):
+    __tablename__ = "PilotSecrets"
+
+    secret_uuid = Column("SecretUUID", Uuid(as_uuid=False), primary_key=True)
+
+    hashed_secret = Column("HashedSecret", BINARY(32))
+    # Global count
+    # Null: Infinite use
+    secret_remaining_use_count = NullColumn(
+        "SecretRemainingUseCount", SmallInteger, default=1
+    )
+    secret_expiration_date = NullColumn("SecretExpirationDate", DateTime(timezone=True))
+    # To authorize only specific pilots to access a secret
+    # The constraint format follows diracx.code.models.PilotSecretConstraints
+    secret_constraints = NullColumn("SecretConstraints", JSON)
+
+    # If a date is set, then it used a secret (acts also like a "PilotUsedSecret" field)
+    pilot_secret_use_date = NullColumn("PilotSecretUseDate", DateTime(timezone=True))
+
+    __table_args__ = (
+        UniqueConstraint("HashedSecret", name="uq_hashed_secret"),
+        Index("HashedSecret", "HashedSecret"),
+    )