feat: We can separate pilot refresh tokens and users refresh tokens

Author: Robin-Van-de-Merghel

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

@@ -99,6 +99,9 @@ def __init__(self, job_id, detail: str = ""):
         )
 
 
+class BadTokenError(DiracError): ...
+
+
 class NotReadyError(DiracError):
     """Tried to access a value which is asynchronously loaded but not yet available."""
 

diracx-core/src/diracx/core/models.py

@@ -7,7 +7,7 @@
 
 import uuid as std_uuid
 from datetime import datetime
-from enum import StrEnum
+from enum import StrEnum, auto
 from typing import Any, Literal, Optional
 
 from pydantic import BaseModel, Field, GetCoreSchemaHandler, GetJsonSchemaHandler
@@ -319,6 +319,7 @@ class PilotFieldsMapping(BaseModel, extra="forbid"):
     AccountingSent: Optional[bool] = None
     CurrentJobID: Optional[int] = None
 
+
 class PilotStatus(StrEnum):
     #: The pilot has been generated and is transferred to a remote site:
     SUBMITTED = "Submitted"
@@ -337,6 +338,7 @@ class PilotStatus(StrEnum):
     #: Cannot get information about the pilot status:
     UNKNOWN = "Unknown"
 
+
 class PilotSecretConstraints(TypedDict, total=False):
     VOs: list[str]  # Authorize only a list of VOs
     PilotStamps: list[str]  # Authorize only a list of stamps
@@ -345,6 +347,13 @@ class PilotSecretConstraints(TypedDict, total=False):
     # We can add constraints here
 
 
+class TokenType(StrEnum):
+    # Pilot token
+    PILOT_TOKEN = auto()
+    # User token
+    USER_TOKEN = auto()
+
+
 class PilotSecretsInfo(BaseModel):
     pilot_secret: str
     pilot_secret_expires_in: int

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

@@ -16,8 +16,8 @@
 )
 from diracx.core.models import (
     PilotFieldsMapping,
-    PilotStatus,
     PilotSecretConstraints,
+    PilotStatus,
     SearchSpec,
     SortSpec,
 )

diracx-db/tests/pilots/test_pilot_management.py

@@ -11,7 +11,6 @@
     PilotFieldsMapping,
     PilotStatus,
 )
-
 from diracx.db.sql.pilots.db import PilotAgentsDB
 
 

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

@@ -23,6 +23,7 @@
     BaseTokenPayload,
     GrantType,
     RefreshTokenPayload,
+    TokenType,
 )
 from diracx.core.properties import SecurityProperty
 from diracx.core.settings import AuthSettings
@@ -159,12 +160,15 @@ async def get_oidc_token_info_from_authorization_flow(
 
 
 async def get_token_info_from_refresh_flow(
-    refresh_token: str, auth_db: AuthDB, settings: AuthSettings
+    refresh_token: str,
+    auth_db: AuthDB,
+    settings: AuthSettings,
+    token_type: TokenType = TokenType.USER_TOKEN,
 ) -> tuple[dict, str, bool, float, bool]:
     """Get OIDC token information from the refresh token DB and check few parameters before returning it."""
     # Decode the refresh token to get the JWT ID
     jti, exp, legacy_exchange = await verify_dirac_refresh_token(
-        refresh_token, settings
+        refresh_token, settings, token_type
     )
 
     # Get some useful user information from the refresh token entry in the DB

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

@@ -15,8 +15,13 @@
 from uuid_utils import UUID
 
 from diracx.core.config.schema import Config
-from diracx.core.exceptions import AuthorizationError, IAMClientError, IAMServerError
-from diracx.core.models import GrantType
+from diracx.core.exceptions import (
+    AuthorizationError,
+    BadTokenError,
+    IAMClientError,
+    IAMServerError,
+)
+from diracx.core.models import GrantType, TokenType
 from diracx.core.properties import SecurityProperty
 from diracx.core.settings import AuthSettings
 
@@ -208,6 +213,7 @@ def read_token(
 async def verify_dirac_refresh_token(
     refresh_token: str,
     settings: AuthSettings,
+    token_type: TokenType = TokenType.USER_TOKEN,
 ) -> tuple[UUID, float, bool]:
     """Verify dirac user token and return a UserInfo class
     Used for each API endpoint.
@@ -216,6 +222,12 @@ async def verify_dirac_refresh_token(
         refresh_token, settings.token_keystore.jwks, settings.token_allowed_algorithms
     )
 
+    if token_type == TokenType.USER_TOKEN and "dirac_policies" not in claims:
+        raise BadTokenError("This is not a user token.")
+
+    if token_type == TokenType.PILOT_TOKEN and "dirac_policies" in claims:
+        raise BadTokenError("This is not a pilot token.")
+
     return (
         UUID(claims["jti"]),
         float(claims["exp"]),

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

@@ -16,6 +16,7 @@
     PilotSecretConstraints,
     PilotSecretsInfo,
     TokenResponse,
+    TokenType,
 )
 from diracx.core.settings import AuthSettings
 from diracx.core.utils import extract_timestamp_from_uuid7, recursive_dict_merge
@@ -359,7 +360,10 @@ async def generate_pilot_tokens(
             _,
             include_refresh_token,
         ) = await get_token_info_from_refresh_flow(
-            refresh_token=refresh_token, auth_db=auth_db, settings=settings
+            refresh_token=refresh_token,
+            auth_db=auth_db,
+            settings=settings,
+            token_type=TokenType.PILOT_TOKEN,
         )
 
         sub = f"{vo}:{pilot_info['sub']}"

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

@@ -10,6 +10,7 @@
 from joserfc.errors import JoseError
 
 from diracx.core.exceptions import (
+    BadTokenError,
     DiracHttpResponseError,
     InvalidCredentialsError,
     PendingAuthorizationError,
@@ -153,7 +154,7 @@ async def get_oidc_token(
             detail=str(e),
             headers={"WWW-Authenticate": "Bearer"},
         ) from e
-    except PermissionError as e:
+    except (BadTokenError, PermissionError) as e:
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,
             detail=str(e),

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

@@ -8,6 +8,7 @@
 
 from diracx.core.exceptions import (
     BadPilotCredentialsError,
+    BadTokenError,
     InvalidCredentialsError,
     PilotNotFoundError,
     SecretHasExpiredError,
@@ -28,8 +29,8 @@
 router = DiracxRouter(require_auth=False)
 
 
-@router.post("/token")
-async def pilot_login(
+@router.post("/secret-exchange")
+async def perform_secret_exchange(
     pilot_db: PilotAgentsDB,
     auth_db: AuthDB,
     pilot_stamp: Annotated[str, Body(description="Stamp used by a pilot to login.")],
@@ -72,7 +73,7 @@ async def pilot_login(
         ) from e
 
 
-@router.post("/refresh-token")
+@router.post("/token")
 async def refresh_pilot_tokens(
     auth_db: AuthDB,
     settings: AuthSettings,
@@ -82,7 +83,8 @@ async def refresh_pilot_tokens(
     ],
     pilot_stamp: Annotated[str, Body(description="Pilot stamp")],
 ) -> TokenResponse:
-    """Endpoint where a pilot can exchange a refresh token for a token."""
+    """Endpoint where *only* pilots can exchange a refresh token for a token."""
+    # Refresh it
     try:
         return await refresh_pilot_token(
             pilot_stamp=pilot_stamp,
@@ -91,7 +93,7 @@ async def refresh_pilot_tokens(
             pilot_db=pilot_db,
             refresh_token=refresh_token,
         )
-    except (InvalidCredentialsError, PilotNotFoundError) as e:
+    except (InvalidCredentialsError, PilotNotFoundError, BadTokenError) as e:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED, detail=str(e)
         ) from e

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

@@ -30,7 +30,7 @@ async def verify_dirac_pilot_access_token(
     settings: AuthSettings,
     authorization: Annotated[str | None, Header()] = None,
 ) -> AuthorizedPilotInfo:
-    """Verify dirac user token and return a UserInfo class
+    """Verify dirac pilot token and return a AuthorizedPilotInfo class
     Used for each API endpoint.
     """
     if not authorization: