fix: Moving pilot auth to auth endpoint, and fixes

Author: Robin-Van-de-Merghel

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

@@ -56,7 +56,7 @@ async def add_pilot_references(
 
     async def increment_pilot_secret_use(
         self,
-        pilot_reference: str,
+        pilot_id: int,
     ) -> None:
 
         #  Prepare the update statement
@@ -65,26 +65,22 @@ async def increment_pilot_secret_use(
             .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)
+            .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_ref=pilot_reference)
+            raise PilotNotFoundError(pilot_id=pilot_id)
 
-    async def verify_pilot_secret(
-        self, pilot_reference: str, pilot_secret: str
-    ) -> None:
+    async def verify_pilot_secret(self, pilot_id: int, pilot_secret: str) -> None:
         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)
+            .where(PilotRegistrations.pilot_id == pilot_id)
         )
 
         # Execute the request
@@ -96,7 +92,7 @@ async def verify_pilot_secret(
             raise AuthorizationError(detail="bad pilot_reference / pilot_secret")
 
         # Increment the count
-        await self.increment_pilot_secret_use(pilot_reference=pilot_reference)
+        await self.increment_pilot_secret_use(pilot_id=pilot_id)
         return
 
     async def register_new_pilot(

diracx-db/tests/pilot_agents/test_pilot_agents_db.py

@@ -73,10 +73,14 @@ async def test_create_pilot_and_verify_secret(pilot_agents_db: PilotAgentsDB):
         assert secret is not None
 
         await pilot_agents_db.verify_pilot_secret(
-            pilot_reference=pilot_ref, pilot_secret=secret
+            pilot_id=new_pilot_id, pilot_secret=secret
         )
 
         with pytest.raises(AuthorizationError):
             await pilot_agents_db.verify_pilot_secret(
-                pilot_reference=pilot_ref, pilot_secret="I love stawberries :)"
+                pilot_id=new_pilot_id, pilot_secret="I love stawberries :)"
+            )
+
+            await pilot_agents_db.verify_pilot_secret(
+                pilot_id=63000, pilot_secret=secret
             )

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

@@ -106,7 +106,7 @@ async def get_oidc_token_info_from_device_flow(
     # TODO: use HTTPException while still respecting the standard format
     # required by the RFC
     if info["Status"] != FlowStatus.READY:
-        # That should never ever happen
+        # That should exnever ever happen
         raise NotImplementedError(f"Unexpected flow status {info['status']!r}")
     return (oidc_token_info, scope)
 
@@ -245,6 +245,39 @@ async def perform_legacy_exchange(
     )
 
 
+def get_verified_preferred_username(
+    config: Config,
+    oidc_token_info: dict,
+    dirac_group: str,
+    properties: set[str],
+    sub: str,
+    vo: str,
+):
+    if user_info := config.Registry[vo].Users.get(sub):
+        preferred_username = user_info.PreferedUsername
+    else:
+        preferred_username = oidc_token_info.get("preferred_username", sub)
+        raise NotImplementedError(
+            "Dynamic registration of users is not yet implemented"
+        )
+
+    # Check that the subject is part of the dirac users
+    if sub not in config.Registry[vo].Groups[dirac_group].Users:
+        raise PermissionError(
+            f"User is not a member of the requested group ({preferred_username}, {dirac_group})"
+        )
+
+    # Check that the user properties are valid
+    allowed_user_properties = get_allowed_user_properties(config, sub, vo)
+    if not properties.issubset(allowed_user_properties):
+        raise PermissionError(
+            f"{' '.join(properties - allowed_user_properties)} are not valid properties "
+            f"for user {preferred_username}, available values: {' '.join(allowed_user_properties)}"
+        )
+
+    return preferred_username
+
+
 async def exchange_token(
     auth_db: AuthDB,
     scope: str,
@@ -255,6 +288,7 @@ async def exchange_token(
     *,
     refresh_token_expire_minutes: int | None = None,
     legacy_exchange: bool = False,
+    pilot_exchange: bool = False,
 ) -> tuple[AccessTokenPayload, RefreshTokenPayload]:
     """Method called to exchange the OIDC token for a DIRAC generated access token."""
     # Extract dirac attributes from the OIDC scope
@@ -265,28 +299,17 @@ async def exchange_token(
 
     # Extract attributes from the OIDC token details
     sub = oidc_token_info["sub"]
-    if user_info := config.Registry[vo].Users.get(sub):
-        preferred_username = user_info.PreferedUsername
-    else:
-        preferred_username = oidc_token_info.get("preferred_username", sub)
-        raise NotImplementedError(
-            "Dynamic registration of users is not yet implemented"
-        )
 
-    # Check that the subject is part of the dirac users
-    if sub not in config.Registry[vo].Groups[dirac_group].Users:
-        raise PermissionError(
-            f"User is not a member of the requested group ({preferred_username}, {dirac_group})"
-        )
+    preferred_username = None
 
-    # Check that the user properties are valid
-    allowed_user_properties = get_allowed_user_properties(config, sub, vo)
-    if not properties.issubset(allowed_user_properties):
-        raise PermissionError(
-            f"{' '.join(properties - allowed_user_properties)} are not valid properties "
-            f"for user {preferred_username}, available values: {' '.join(allowed_user_properties)}"
+    if not pilot_exchange:
+        preferred_username = get_verified_preferred_username(
+            config, oidc_token_info, dirac_group, properties, sub, vo
         )
 
+    else:
+        preferred_username = oidc_token_info["pilot_reference"]
+
     # Merge the VO with the subject to get a unique DIRAC sub
     sub = f"{vo}:{sub}"
 

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

@@ -5,6 +5,7 @@
 from .authorize_code_flow import router as authorize_code_flow_router
 from .device_flow import router as device_flow_router
 from .management import router as management_router
+from .pilot_auth import router as pilot_auth_router
 from .token import router as token_router
 from .utils import has_properties
 
@@ -13,5 +14,6 @@
 router.include_router(management_router)
 router.include_router(authorize_code_flow_router)
 router.include_router(token_router)
+router.include_router(pilot_auth_router)
 
 __all__ = ["has_properties", "verify_dirac_access_token"]

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

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from fastapi import HTTPException, status
+
+from diracx.core.exceptions import AuthorizationError
+from diracx.logic.auth.token import create_token, exchange_token
+
+from ..dependencies import (
+    AuthDB,
+    AuthSettings,
+    AvailableSecurityProperties,
+    Config,
+    PilotAgentsDB,
+)
+from ..fastapi_classes import DiracxRouter
+
+router = DiracxRouter(require_auth=False)
+
+
+@router.post("/pilot-login")
+async def pilot_login(
+    pilot_db: PilotAgentsDB,
+    auth_db: AuthDB,
+    pilot_id: int,
+    pilot_secret: str,
+    config: Config,
+    settings: AuthSettings,
+    available_properties: AvailableSecurityProperties,
+):
+    """Endpoint without policy, the pilot uses only its secret."""
+    try:
+        await pilot_db.verify_pilot_secret(pilot_id=pilot_id, pilot_secret=pilot_secret)
+    except AuthorizationError as e:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED, detail=e.detail
+        ) from e
+
+    pilot = await pilot_db.get_pilot_by_id(pilot_id=pilot_id)
+
+    pilot_info = {
+        "pilot_reference": pilot["PilotJobReference"],
+        "sub": pilot["PilotJobReference"],
+    }
+
+    # return pilot_info
+    access_token, refresh_token = await exchange_token(
+        auth_db=auth_db,
+        scope="vo:diracAdmin",
+        oidc_token_info=pilot_info,
+        config=config,
+        settings=settings,
+        available_properties=available_properties,
+        pilot_exchange=True,
+    )
+
+    return [create_token(access_token, settings), create_token(refresh_token, settings)]

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

@@ -3,9 +3,9 @@
 import logging
 
 from ..fastapi_classes import DiracxRouter
-from .auth import router as auth_router
+from .debug import router as debug_router
 
 logger = logging.getLogger(__name__)
 
-router = DiracxRouter(require_auth=False)
-router.include_router(auth_router)
+router = DiracxRouter()
+router.include_router(debug_router)

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

@@ -7,7 +7,7 @@
 from fastapi import Depends, HTTPException, status
 
 # TODO:  DEBUG
-from diracx.core.properties import GENERIC_PILOT, LIMITED_DELEGATION, NORMAL_USER
+from diracx.core.properties import GENERIC_PILOT, LIMITED_DELEGATION
 from diracx.db.sql import PilotAgentsDB
 from diracx.routers.access_policies import BaseAccessPolicy
 from diracx.routers.utils.users import AuthorizedUserInfo
@@ -33,22 +33,15 @@ async def policy(
         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"
         )

diracx-routers/src/diracx/routers/pilots/auth.py renamed to diracx-routers/src/diracx/routers/pilots/debug.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
 from os import getenv
+from typing import Annotated
 
-from fastapi import HTTPException, status
+from fastapi import (
+    Depends,
+    HTTPException,
+    status,
+)
 
 from diracx.core.exceptions import (
-    AuthorizationError,
     PilotAlreadyExistsError,
     PilotNotFoundError,
 )
@@ -14,31 +18,25 @@
     PilotAgentsDB,
 )
 from ..fastapi_classes import DiracxRouter
+from ..utils.users import AuthorizedUserInfo, verify_dirac_access_token
+from .access_policies import RegisteredPilotAccessPolicyCallable
 
-router = DiracxRouter(require_auth=False)
-
-
-@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.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
-
-    # TODO: Returns a JWT
-    return {"ref": pilot_reference}
-
+router = DiracxRouter()
 
 # Debug only route
 # Keep it?
 # TODO: Add to the env?
 if getenv("production") is None:
 
+    @router.get("/info")
+    async def get_pilot_info(
+        check_permissions: RegisteredPilotAccessPolicyCallable,
+        user_info: Annotated[AuthorizedUserInfo, Depends(verify_dirac_access_token)],
+    ):
+        await check_permissions()
+
+        return user_info
+
     @router.post("/register-pilot")
     async def add_new_pilot(
         pilot_db: PilotAgentsDB,