feat: Better policy, to support also jobs

Author: Robin VAN DE MERGHEL

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

@@ -7,6 +7,7 @@
 from fastapi import Depends, HTTPException, status
 
 from diracx.core.properties import SERVICE_ADMINISTRATOR
+from diracx.db.sql.job.db import JobDB
 from diracx.db.sql.pilots.db import PilotAgentsDB
 from diracx.logic.pilots.query import get_pilots_by_stamp
 from diracx.routers.access_policies import BaseAccessPolicy
@@ -35,26 +36,50 @@ async def policy(
         action: ActionType | None = None,
         pilot_db: PilotAgentsDB | None = None,
         pilot_stamps: list[str] | None = None,
+        job_db: JobDB | None = None,
+        job_ids: list[int] | None = None,
     ):
         assert action, "action is a mandatory parameter"
 
         # Users can query
         # NOTE: Add into queries a VO constraint
-        if action == ActionType.READ_PILOT_FIELDS:
-            return
+        # To manage pilots, user have to be an admin
+        if (
+            action == ActionType.MANAGE_PILOTS
+            and SERVICE_ADMINISTRATOR not in user_info.properties
+        ):
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="You don't have the permission to manage pilots.",
+            )
 
-        # If we want to modify pilots, we allow only admins
-        # TODO: See if we add other types of admins
-        if SERVICE_ADMINISTRATOR in user_info.properties:
-            # If we don't provide pilot_db and pilot_stamps, we accept directly
-            # This is for example when we submit pilots, we use the user VO, so no need to verify
-            if not (pilot_db and pilot_stamps):
-                return
+        #
+        # Additional checks if job_ids or pilot_stamps are provided
+        #
 
-            # Else, check its VO
-            assert pilot_db, "PilotDB is needed to determine pilot VO."
-            assert pilot_stamps, "PilotStamps are needed to determine pilot VO."
+        # First, if job_ids are provided, we check who is the owner
+        if job_db and job_ids:
+            job_owners = await job_db.summary(
+                ["Owner", "VO"],
+                [{"parameter": "JobID", "operator": "in", "values": job_ids}],
+            )
+
+            expected_owner = {
+                "Owner": user_info.preferred_username,
+                "VO": user_info.vo,
+                "count": len(set(job_ids)),
+            }
+            # All the jobs belong to the user doing the query
+            # and all of them are present
+            if not job_owners == [expected_owner]:
+                raise HTTPException(
+                    status_code=status.HTTP_403_FORBIDDEN,
+                    detail="You don't have the rights to modify a pilot.",
+                )
 
+        # This is for example when we submit pilots, we use the user VO, so no need to verify
+        if pilot_db and pilot_stamps:
+            # Else, check its VO
             pilots = await get_pilots_by_stamp(
                 pilot_db=pilot_db,
                 pilot_stamps=pilot_stamps,
@@ -74,13 +99,6 @@ async def policy(
                     detail="You don't have access to all pilots.",
                 )
 
-            return
-
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail="You don't have the rights to modify a pilot.",
-        )
-
 
 CheckPilotManagementPolicyCallable = Annotated[
     Callable, Depends(PilotManagementAccessPolicy.check)

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

@@ -28,7 +28,7 @@
 from diracx.logic.pilots.query import get_pilot_ids_by_job_id
 from diracx.routers.utils.users import AuthorizedUserInfo, verify_dirac_access_token
 
-from ..dependencies import PilotAgentsDB
+from ..dependencies import JobDB, PilotAgentsDB
 from ..fastapi_classes import DiracxRouter
 from .access_policies import (
     ActionType,
@@ -203,6 +203,7 @@ async def update_pilot_fields(
 @router.get("/jobs")
 async def get_pilot_jobs(
     pilot_db: PilotAgentsDB,
+    job_db: JobDB,
     check_permissions: CheckPilotManagementPolicyCallable,
     pilot_stamp: Annotated[
         str | None, Query(description="The stamp of the pilot.")
@@ -211,14 +212,23 @@ async def get_pilot_jobs(
 ) -> list[int]:
     """Endpoint only for admins, to get jobs of a pilot."""
     if pilot_stamp:
-        await check_permissions(action=ActionType.READ_PILOT_FIELDS)
+        # Check VO
+        await check_permissions(
+            action=ActionType.READ_PILOT_FIELDS,
+            pilot_db=pilot_db,
+            pilot_stamps=[pilot_stamp],
+        )
 
         return await get_pilot_jobs_ids_by_stamp(
             pilot_db=pilot_db,
             pilot_stamp=pilot_stamp,
         )
     elif job_id:
-        # FIXME: Add some policy, verify that it's the user's job?
+        # Check job owner
+        await check_permissions(
+            action=ActionType.READ_PILOT_FIELDS, job_db=job_db, job_ids=[job_id]
+        )
+
         return await get_pilot_ids_by_job_id(pilot_db=pilot_db, job_id=job_id)
 
     raise HTTPException(
@@ -230,6 +240,7 @@ async def get_pilot_jobs(
 @router.patch("/jobs", status_code=HTTPStatus.NO_CONTENT)
 async def add_jobs_to_pilot(
     pilot_db: PilotAgentsDB,
+    job_db: JobDB,
     pilot_stamp: Annotated[str, Body(description="The stamp of the pilot.")],
     job_ids: Annotated[
         list[int], Body(description="The jobs we want to add to the pilot.")
@@ -238,9 +249,12 @@ async def add_jobs_to_pilot(
 ):
     """Endpoint only for admins, to associate a pilot with a job."""
     await check_permissions(
-        action=ActionType.MANAGE_PILOTS, pilot_db=pilot_db, pilot_stamps=[pilot_stamp]
+        action=ActionType.MANAGE_PILOTS,
+        pilot_db=pilot_db,
+        pilot_stamps=[pilot_stamp],
+        job_db=job_db,
+        job_ids=job_ids,
     )
-    # FIXME: Also verify job_ids
 
     try:
         await add_jobs_to_pilot_bl(