feat: Adding pilot registrations

Author: Robin Van de Merghel

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

@@ -101,3 +101,41 @@ def __init__(self, job_id, detail: str | None = None):
         super().__init__(
             f"Error concerning job {job_id}" + (": {detail} " if detail else "")
         )
+
+
+class PilotNotFoundError(Exception):
+    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__(
+            "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

@@ -2,10 +2,17 @@
 
 from datetime import datetime, timezone
 
-from sqlalchemy import insert
+from sqlalchemy import DateTime, insert, select, update
+from sqlalchemy.exc import IntegrityError, NoResultFound
+
+from diracx.core.exceptions import (
+    AuthorizationError,
+    PilotAlreadyExistsError,
+    PilotNotFoundError,
+)
 
 from ..utils import BaseSQLDB
-from .schema import PilotAgents, PilotAgentsDBBase
+from .schema import PilotAgents, PilotAgentsDBBase, PilotRegistrations
 
 
 class PilotAgentsDB(BaseSQLDB):
@@ -44,3 +51,112 @@ async def add_pilot_references(
         stmt = insert(PilotAgents).values(values)
         await self.conn.execute(stmt)
         return
+
+    async def increment_pilot_secret_use(
+        self,
+        pilot_id: int,
+    ) -> None:
+
+        #  Prepare the update statement
+        stmt = (
+            update(PilotRegistrations)
+            .values(
+                pilot_secret_use_count=PilotRegistrations.pilot_secret_use_count + 1
+            )
+            .where(PilotRegistrations.pilot_id == pilot_id)
+        )
+
+        # Execute the update using the connection
+        res = await self.conn.execute(stmt)
+
+        if res.rowcount == 0:
+            raise PilotNotFoundError(pilot_id=pilot_id)
+
+    async def verify_pilot_secret(
+        self, pilot_job_reference: str, pilot_hashed_secret: str
+    ) -> None:
+
+        try:
+            pilot = await self.get_pilot_by_reference(pilot_job_reference)
+        except NoResultFound as e:
+            raise PilotNotFoundError(pilot_ref=pilot_job_reference) from e
+
+        pilot_id = pilot["PilotID"]
+
+        stmt = (
+            select(PilotRegistrations)
+            .where(PilotRegistrations.pilot_hashed_secret == pilot_hashed_secret)
+            .where(PilotRegistrations.pilot_id == pilot_id)
+        )
+
+        # Execute the request
+        res = await self.conn.execute(stmt)
+
+        result = res.fetchone()
+
+        if result is None:
+            raise AuthorizationError(detail="bad pilot_id / pilot_secret")
+
+        # Increment the count
+        await self.increment_pilot_secret_use(pilot_id=pilot_id)
+
+    async def register_new_pilot(
+        self,
+        vo: str,
+        pilot_job_reference: str,
+        pilot_stamp: str,
+        grid_type: str = "DIRAC",
+        submission_time: DateTime | None = None,  # ?
+        last_update_time: DateTime | None = None,  # = now?
+    ) -> int | None:
+        stmt = insert(PilotAgents).values(
+            vo=vo,
+            submission_time=submission_time,
+            last_update_time=last_update_time,
+            pilot_job_reference=pilot_job_reference,
+            grid_type=grid_type,
+            pilot_stamp=pilot_stamp,
+        )
+
+        # 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, pilot_hashed_secret: str):
+
+        stmt = insert(PilotRegistrations).values(
+            pilot_id=pilot_id, pilot_hashed_secret=pilot_hashed_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
+
+    async def fetch_all_pilots(self):
+        stmt = select(PilotRegistrations).with_for_update()
+        result = await self.conn.execute(stmt)
+
+        # Convert results into a dictionary
+        pilots = [dict(row._mapping) for row in result]
+
+        return pilots
+
+    async def get_pilot_by_reference(self, pilot_ref: str):
+        stmt = (
+            select(PilotAgents)
+            .with_for_update()
+            .where(PilotAgents.pilot_job_reference == pilot_ref)
+        )
+
+        # We assume it is unique...
+        return dict((await self.conn.execute(stmt)).one()._mapping)

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
@@ -58,3 +51,16 @@ class PilotOutput(PilotAgentsDBBase):
     pilot_id = Column("PilotID", Integer, primary_key=True)
     std_output = Column("StdOutput", Text)
     std_error = Column("StdError", Text)
+
+
+class PilotRegistrations(PilotAgentsDBBase):
+    __tablename__ = "PilotRegistrations"
+
+    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

@@ -1,12 +1,15 @@
 from __future__ import annotations
 
 import pytest
+from sqlalchemy.exc import NoResultFound
 
+from diracx.core.exceptions import AuthorizationError
 from diracx.db.sql.pilot_agents.db import PilotAgentsDB
+from diracx.db.sql.utils.functions import hash
 
 
 @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 +32,66 @@ 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:
+        pilot_reference = "pilot-reference-test"
+        await pilot_agents_db.register_new_pilot(
+            vo="pilot-vo",
+            pilot_job_reference=pilot_reference,
+            pilot_stamp="pilot-stamp",
+            grid_type="grid-type",
+        )
+
+        res = await pilot_agents_db.get_pilot_by_reference(pilot_ref=pilot_reference)
+
+        with pytest.raises(NoResultFound):
+            await pilot_agents_db.get_pilot_by_reference("I am a fake ref")
+
+        # Set values
+        assert res["VO"] == "pilot-vo"
+        assert res["PilotJobReference"] == pilot_reference
+        assert res["PilotStamp"] == "pilot-stamp"
+        assert res["GridType"] == "grid-type"
+
+        # Default values
+        assert res["BenchMark"] == 0.0
+        assert res["Status"] == "Unknown"
+
+
+async def test_create_pilot_and_verify_secret(pilot_agents_db: PilotAgentsDB):
+
+    async with pilot_agents_db as pilot_agents_db:
+        pilot_reference = "pilot-reference-test"
+        pilot_id = await pilot_agents_db.register_new_pilot(
+            vo="pilot-vo",
+            pilot_job_reference=pilot_reference,
+            pilot_stamp="pilot-stamp",
+        )
+
+        secret = "AW0nd3rfulS3cr3t"
+        pilot_hashed_secret = hash(secret)
+
+        # Add creds
+        await pilot_agents_db.add_pilot_credentials(
+            pilot_id=pilot_id, pilot_hashed_secret=pilot_hashed_secret
+        )
+
+        assert secret is not None
+
+        await pilot_agents_db.verify_pilot_secret(
+            pilot_job_reference=pilot_reference, pilot_hashed_secret=pilot_hashed_secret
+        )
+
+        with pytest.raises(AuthorizationError):
+            await pilot_agents_db.verify_pilot_secret(
+                pilot_job_reference=pilot_reference,
+                pilot_hashed_secret=hash("I love stawberries :)"),
+            )
+
+            await pilot_agents_db.verify_pilot_secret(
+                pilot_job_reference="I am a spider",
+                pilot_hashed_secret=pilot_hashed_secret,
+            )

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

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from secrets import token_hex
+
+from diracx.db.sql import PilotAgentsDB
+
+# TODO: Move this hash function in diracx-logic, and rename it
+from diracx.db.sql.utils.functions import hash
+
+
+def generate_pilot_secret() -> str:
+    # Can change with time
+    return token_hex(32)
+
+
+async def add_pilot_credentials(pilot_id: int, pilot_db: PilotAgentsDB) -> str:
+
+    # Get a random string
+    # Can be customized
+    random_secret = generate_pilot_secret()
+
+    hashed_secret = hash(random_secret)
+
+    await pilot_db.add_pilot_credentials(
+        pilot_id=pilot_id, pilot_hashed_secret=hashed_secret
+    )
+
+    return random_secret
+
+
+def generate_pilot_scope(pilot: dict) -> str:
+    return f"vo:{pilot['VO']}"
+
+
+async def try_login(
+    pilot_reference: str, pilot_db: PilotAgentsDB, pilot_secret: str
+) -> None:
+
+    hashed_secret = hash(pilot_secret)
+
+    await pilot_db.verify_pilot_secret(
+        pilot_hashed_secret=hashed_secret, pilot_job_reference=pilot_reference
+    )