feat: Adding pilot secret creation

Author: Robin-Van-de-Merghel

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

@@ -104,9 +104,38 @@ def __init__(self, job_id, detail: str | None = None):
 
 
 class PilotNotFoundError(Exception):
-    def __init__(self, pilot_ref: str, detail: str | None = None):
+    def __init__(
+        self,
+        pilot_ref: str | None = None,
+        pilot_id: int | None = None,
+        detail: str | None = None,
+    ):
         self.pilot_ref = pilot_ref
+        self.pilot_id = pilot_id
         self.detail = detail
         super().__init__(
-            f"Pilot {pilot_ref} not found" + (": {detail} " if detail else "")
+            "Pilot "
+            + (f"(Ref: {pilot_ref})" if pilot_ref else "")
+            + (f" (ID: {str(pilot_id)})" if pilot_id is not None else "")
+            + " not found"
+            + (f": {detail}" if detail else "")
+        )
+
+
+class PilotAlreadyExistsError(Exception):
+    def __init__(
+        self,
+        pilot_ref: str | None = None,  # Changed to str based on the format
+        pilot_id: int | None = None,
+        detail: str | None = None,
+    ):
+        self.pilot_ref = pilot_ref
+        self.pilot_id = pilot_id
+        self.detail = detail
+        super().__init__(
+            "Pilot "
+            + (f"(Ref: {pilot_ref})" if pilot_ref else "")
+            + (f" (ID: {str(pilot_id)})" if pilot_id is not None else "")
+            + " already exists"
+            + (f": {detail}" if detail else "")
         )

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

@@ -1,14 +1,17 @@
 from __future__ import annotations
 
-import hashlib
 from datetime import datetime, timezone
+from os import urandom
 
 from sqlalchemy import insert, select, update
+from sqlalchemy.exc import IntegrityError
 
 from diracx.core.exceptions import (
     AuthorizationError,
+    PilotAlreadyExistsError,
     PilotNotFoundError,
 )
+from diracx.db.sql.utils.functions import hash
 
 from ..utils import BaseSQLDB
 from .schema import PilotAgents, PilotAgentsDBBase, PilotRegistrations
@@ -51,6 +54,28 @@ async def add_pilot_references(
         await self.conn.execute(stmt)
         return
 
+    async def register_pilot(self, pilot_reference: str, pilot_secret: str):
+        hashed_secret = hash(pilot_secret)
+
+        stmt = (
+            select(PilotRegistrations)
+            .join(PilotAgents, PilotRegistrations.pilot_id == PilotAgents.pilot_id)
+            .where(PilotRegistrations.pilot_hashed_secret == hashed_secret)
+            .where(PilotAgents.pilot_job_reference == pilot_reference)
+        )
+
+        # Execute the request
+        res = await self.conn.execute(stmt)
+
+        result = res.fetchone()
+
+        if result is None:
+            raise AuthorizationError(detail="bad pilot_reference / pilot_secret")
+
+        # Increment the count
+        await self.increment_pilot_secret_use(pilot_reference=pilot_reference)
+        return
+
     async def increment_pilot_secret_use(
         self,
         pilot_reference: str,
@@ -72,10 +97,10 @@ async def increment_pilot_secret_use(
         if res.rowcount == 0:
             raise PilotNotFoundError(pilot_ref=pilot_reference)
 
-    async def register_pilot(self, pilot_reference: str, pilot_secret: str):
-        hashed_secret = hashlib.sha256(pilot_secret.encode()).hexdigest()
-
-        print("hash", hashed_secret)
+    async def verify_pilot_secret(
+        self, pilot_reference: str, pilot_secret: str
+    ) -> None:
+        hashed_secret = hash(pilot_secret)
 
         stmt = (
             select(PilotRegistrations)
@@ -95,3 +120,85 @@ async def register_pilot(self, pilot_reference: str, pilot_secret: str):
         # Increment the count
         await self.increment_pilot_secret_use(pilot_reference=pilot_reference)
         return
+
+    async def register_new_pilot(
+        self,
+        vo: str,
+        initial_job_id: int = 0,
+        current_job_id: int = 0,
+        benchmark: float = 0.0,
+        pilot_job_reference: str = "Unknown",
+        pilot_stamp: str = "",
+        status: str = "Unknown",
+        status_reason: str = "Unknown",
+        queue: str = "Unknown",
+        grid_site: str = "Unknown",
+        destination_site: str = "NotAssigned",
+        grid_type: str = "LCG",
+        submission_time: str | None = None,  # ?
+        last_update_time: str | None = None,  # = now?
+        accounting_sent: bool = False,
+    ) -> int | None:
+        stmt = insert(PilotAgents).values(
+            initial_job_id=initial_job_id,
+            current_job_id=current_job_id,
+            pilot_job_reference=pilot_job_reference,
+            pilot_stamp=pilot_stamp,
+            destination_site=destination_site,
+            queue=queue,
+            grid_site=grid_site,
+            vo=vo,
+            grid_type=grid_type,
+            benchmark=benchmark,
+            submission_time=submission_time,
+            last_update_time=last_update_time,
+            status=status,
+            status_reason=status_reason,
+            accounting_sent=accounting_sent,
+        )
+
+        # Execute the request
+        res = await self.conn.execute(stmt)
+
+        new_pilot_id = res.inserted_primary_key
+
+        # Returns the new pilot ID
+        return int(new_pilot_id[0]) if new_pilot_id else None
+
+    async def add_pilot_credentials(self, pilot_id: int) -> str:
+
+        # Get a random string
+        # Can be customized
+        random_secret = urandom(30).hex()
+
+        hashed_random_secret = hash(random_secret)
+
+        stmt = insert(PilotRegistrations).values(
+            pilot_id=pilot_id, pilot_hashed_secret=hashed_random_secret
+        )
+
+        try:
+            await self.conn.execute(stmt)
+        except IntegrityError as e:
+            if "foreign key" in str(e.orig).lower():
+                raise PilotNotFoundError(pilot_id=pilot_id) from e
+            if "duplicate entry" in str(e.orig).lower():
+                raise PilotAlreadyExistsError(
+                    pilot_id=pilot_id, detail="this pilot has already credentials"
+                ) from e
+
+        return random_secret
+
+    async def get_pilots(self):
+        """Récupère tous les pilotes et les retourne sous forme de dictionnaires.
+
+        :raises: NoResultFound
+        """
+        # La clause with_for_update empêche que le jeton soit récupéré plusieurs fois simultanément
+        stmt = select(PilotRegistrations).with_for_update()
+        result = await self.conn.execute(stmt)
+
+        # Convertir les résultats en dictionnaires
+        pilots = [dict(row._mapping) for row in result]
+
+        return pilots

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

@@ -1,13 +1,6 @@
 from __future__ import annotations
 
-from sqlalchemy import (
-    DateTime,
-    Double,
-    Index,
-    Integer,
-    String,
-    Text,
-)
+from sqlalchemy import DateTime, Double, ForeignKey, Index, Integer, String, Text
 from sqlalchemy.orm import declarative_base
 
 from ..utils import Column, EnumBackedBool, NullColumn
@@ -63,6 +56,11 @@ class PilotOutput(PilotAgentsDBBase):
 class PilotRegistrations(PilotAgentsDBBase):
     __tablename__ = "PilotRegistrations"
 
-    pilot_id = Column("PilotID", Integer, primary_key=True)
+    pilot_id = Column(
+        "PilotID",
+        Integer,
+        ForeignKey("PilotAgents.PilotID", ondelete="CASCADE"),
+        primary_key=True,
+    )
     pilot_hashed_secret = Column("PilotHashedSecret", String(64))
     pilot_secret_use_count = Column("PilotSecretUseCount", Integer, default=0)

diracx-db/tests/pilot_agents/test_pilot_agents_db.py

@@ -6,7 +6,7 @@
 
 
 @pytest.fixture
-async def pilot_agents_db(tmp_path) -> PilotAgentsDB:
+async def pilot_agents_db(tmp_path):
     agents_db = PilotAgentsDB("sqlite+aiosqlite:///:memory:")
     async with agents_db.engine_context():
         async with agents_db.engine.begin() as conn:
@@ -29,3 +29,24 @@ async def test_insert_and_select(pilot_agents_db: PilotAgentsDB):
         await pilot_agents_db.add_pilot_references(
             refs, "test_vo", grid_type="DIRAC", pilot_stamps=None
         )
+
+
+async def test_insert_and_select_single(pilot_agents_db: PilotAgentsDB):
+
+    async with pilot_agents_db as pilot_agents_db:
+        # Add a pilot reference
+        new_pilot_id = await pilot_agents_db.register_new_pilot(
+            vo="pilot-vo", pilot_job_reference="pilot-ref"
+        )
+
+        res = await pilot_agents_db.get_pilot_by_id(new_pilot_id)
+
+        # Set values
+        assert res["PilotID"] == new_pilot_id
+        assert res["VO"] == "pilot-vo"
+        assert res["PilotJobReference"] == "pilot-ref"
+
+        # Default values
+        assert res["PilotStamp"] == ""
+        assert res["BenchMark"] == 0.0
+        assert res["Status"] == "Unknown"

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

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
+from os import getenv
+
 from fastapi import HTTPException, status
 
 from diracx.core.exceptions import (
     AuthorizationError,
+    PilotAlreadyExistsError,
+    PilotNotFoundError,
 )
 
 from ..dependencies import (
@@ -14,16 +18,79 @@
 router = DiracxRouter(require_auth=False)
 
 
-@router.post("/register")
-async def search(pilot_db: PilotAgentsDB, pilot_reference: str, pilot_secret: str):
+@router.post("/pilot-login")
+async def pilot_login(pilot_db: PilotAgentsDB, pilot_reference: str, pilot_secret: str):
     """Endpoint without policy, the pilot uses only its secret."""
     try:
-        await pilot_db.register_pilot(
+        await pilot_db.verify_pilot_secret(
             pilot_reference=pilot_reference, pilot_secret=pilot_secret
         )
     except AuthorizationError as e:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED, detail=e.detail
         ) from e
 
-    return pilot_reference
+    # TODO: Returns a JWT
+    return {"ref": pilot_reference}
+
+
+# Debug only route
+# Keep it?
+# TODO: Add to the env?
+if getenv("production") is None:
+
+    @router.post("/register-pilot")
+    async def add_new_pilot(
+        pilot_db: PilotAgentsDB,
+        vo: str = "default-vo",
+        initial_job_id: int = 0,
+        current_job_id: int = 0,
+        benchmark: float = 0.0,
+        pilot_job_reference: str = "Unknown",
+        pilot_stamp: str = "",
+        status: str = "Unknown",
+        status_reason: str = "Unknown",
+        queue: str = "Unknown",
+        grid_site: str = "Unknown",
+        destination_site: str = "NotAssigned",
+        grid_type: str = "LCG",
+        submission_time: str | None = None,  # ?
+        last_update_time: str | None = None,  # = now?
+        accounting_sent: bool = False,
+    ):
+        return {
+            "id": await pilot_db.register_new_pilot(
+                initial_job_id=initial_job_id,
+                current_job_id=current_job_id,
+                pilot_job_reference=pilot_job_reference,
+                pilot_stamp=pilot_stamp,
+                destination_site=destination_site,
+                queue=queue,
+                grid_site=grid_site,
+                vo=vo,
+                grid_type=grid_type,
+                benchmark=benchmark,
+                submission_time=submission_time,
+                last_update_time=last_update_time,
+                status=status,
+                status_reason=status_reason,
+                accounting_sent=accounting_sent,
+            )
+        }
+
+    @router.post("/add-credentials")
+    async def add_credentials(pilot_db: PilotAgentsDB, pilot_id: int):
+        try:
+            return {"secret": await pilot_db.add_pilot_credentials(pilot_id=pilot_id)}
+        except PilotNotFoundError as e:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND, detail=e.detail
+            ) from e
+        except PilotAlreadyExistsError as e:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT, detail=e.detail
+            ) from e
+
+    @router.post("/get-pilots")
+    async def get_pilots(pilot_db: PilotAgentsDB):
+        return await pilot_db.get_pilots()