feat: Pilot can exchange against nothing a secret

Author: Robin-Van-de-Merghel

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

@@ -101,3 +101,12 @@ 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, detail: str | None = None):
+        self.pilot_ref = pilot_ref
+        self.detail = detail
+        super().__init__(
+            f"Pilot {pilot_ref} not found" + (": {detail} " if detail else "")
+        )

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

@@ -1,11 +1,17 @@
 from __future__ import annotations
 
+import hashlib
 from datetime import datetime, timezone
 
-from sqlalchemy import insert
+from sqlalchemy import insert, select, update
+
+from diracx.core.exceptions import (
+    AuthorizationError,
+    PilotNotFoundError,
+)
 
 from ..utils import BaseSQLDB
-from .schema import PilotAgents, PilotAgentsDBBase
+from .schema import PilotAgents, PilotAgentsDBBase, PilotRegistrations
 
 
 class PilotAgentsDB(BaseSQLDB):
@@ -44,3 +50,48 @@ async def add_pilot_references(
         stmt = insert(PilotAgents).values(values)
         await self.conn.execute(stmt)
         return
+
+    async def increment_pilot_secret_use(
+        self,
+        pilot_reference: str,
+    ) -> None:
+
+        #  Prepare the update statement
+        stmt = (
+            update(PilotRegistrations)
+            .values(
+                pilot_secret_use_count=PilotRegistrations.pilot_secret_use_count + 1
+            )
+            .where(PilotRegistrations.pilot_id == PilotAgents.pilot_id)
+            .where(PilotAgents.pilot_job_reference == pilot_reference)
+        )
+
+        # Execute the update using the connection
+        res = await self.conn.execute(stmt)
+
+        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)
+
+        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

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

@@ -58,3 +58,11 @@ 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, primary_key=True)
+    pilot_hashed_secret = Column("PilotHashedSecret", String(64))
+    pilot_secret_use_count = Column("PilotSecretUseCount", Integer, default=0)

diracx-routers/pyproject.toml

@@ -44,11 +44,13 @@ types = [
 jobs = "diracx.routers.jobs:router"
 config = "diracx.routers.configuration:router"
 auth = "diracx.routers.auth:router"
+pilots = "diracx.routers.pilots:router"
 ".well-known" = "diracx.routers.auth.well_known:router"
 
 [project.entry-points."diracx.access_policies"]
 WMSAccessPolicy = "diracx.routers.jobs.access_policies:WMSAccessPolicy"
 SandboxAccessPolicy = "diracx.routers.jobs.access_policies:SandboxAccessPolicy"
+RegisteredPilotAccessPolicy = "diracx.routers.pilots.access_policies:RegisteredPilotAccessPolicy"
 
 # Minimum version of the client supported
 [project.entry-points."diracx.min_client_version"]

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

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import logging
+
+from ..fastapi_classes import DiracxRouter
+from .auth import router as auth_router
+
+logger = logging.getLogger(__name__)
+
+router = DiracxRouter(require_auth=False)
+router.include_router(auth_router)

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

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from enum import StrEnum, auto
+from typing import Annotated
+
+from fastapi import Depends, HTTPException, status
+
+# TODO:  DEBUG
+from diracx.core.properties import GENERIC_PILOT, LIMITED_DELEGATION, NORMAL_USER
+from diracx.db.sql import PilotAgentsDB
+from diracx.routers.access_policies import BaseAccessPolicy
+from diracx.routers.utils.users import AuthorizedUserInfo
+
+
+class ActionType(StrEnum):
+    #: Create a job or a sandbox
+    CREATE = auto()
+    #: Check job status, download a sandbox
+    READ = auto()
+    #: delete, kill, remove, set status, etc of a job
+    #: delete or assign a sandbox
+    MANAGE = auto()
+    #: Search
+    QUERY = auto()
+
+
+class RegisteredPilotAccessPolicy(BaseAccessPolicy):
+
+    @staticmethod
+    async def policy(
+        policy_name: str,
+        pilot_info: AuthorizedUserInfo,
+        /,
+        *,
+        action: ActionType | None = None,
+        pilot_db: PilotAgentsDB | None = None,
+    ):
+
+        assert pilot_db, "pilot_db is a mandatory parameter"
+
+        if GENERIC_PILOT in pilot_info.properties:
+            return
+
+        if LIMITED_DELEGATION in pilot_info.properties:
+            return
+
+        # TODO: DEBUG
+        if NORMAL_USER in pilot_info.properties:
+            return
+
+        raise HTTPException(
+            status.HTTP_401_UNAUTHORIZED, "you don't have the right properties"
+        )
+        return
+
+
+RegisteredPilotAccessPolicyCallable = Annotated[
+    Callable, Depends(RegisteredPilotAccessPolicy.check)
+]

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

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from fastapi import HTTPException, status
+
+from diracx.core.exceptions import (
+    AuthorizationError,
+)
+
+from ..dependencies import (
+    PilotAgentsDB,
+)
+from ..fastapi_classes import DiracxRouter
+
+router = DiracxRouter(require_auth=False)
+
+
+@router.post("/register")
+async def search(pilot_db: PilotAgentsDB, pilot_reference: str, pilot_secret: str):
+    """Endpoint without policy, the pilot uses only its secret."""
+    try:
+        await pilot_db.register_pilot(
+            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