feat: Adding duration to secrets

Author: Robin-Van-de-Merghel

diracx-core/src/diracx/core/settings.py

@@ -161,6 +161,7 @@ class AuthSettings(ServiceSettingsBase):
     allowed_redirects: list[str] = []
     device_flow_expiration_seconds: int = 600
     authorization_flow_expiration_seconds: int = 300
+    pilot_secret_expire_seconds: int = 600
 
     # State key is used to encrypt/decrypt the state dict passed to the IAM
     state_key: FernetKey

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

@@ -2,7 +2,7 @@
 
 from datetime import datetime, timezone
 
-from sqlalchemy import insert, select, update
+from sqlalchemy import DateTime, insert, select, update
 from sqlalchemy.exc import IntegrityError, NoResultFound
 
 from diracx.core.exceptions import (
@@ -85,8 +85,13 @@ async def verify_pilot_secret(
 
         stmt = (
             select(PilotRegistrations)
+            .with_for_update()
             .where(PilotRegistrations.pilot_hashed_secret == pilot_hashed_secret)
             .where(PilotRegistrations.pilot_id == pilot_id)
+            .where(
+                PilotRegistrations.pilot_secret_expiration_date
+                > datetime.now(tz=timezone.utc)
+            )
         )
 
         # Execute the request
@@ -95,12 +100,16 @@ async def verify_pilot_secret(
         result = res.fetchone()
 
         if result is None:
-            raise AuthorizationError(detail="bad pilot_id / pilot_secret")
+            raise AuthorizationError(
+                detail="bad pilot_id / pilot_secret or secret has expired"
+            )
 
         # Increment the count
         await self.increment_pilot_secret_use(pilot_id=pilot_id)
 
-    async def add_pilot_credentials(self, pilot_id: int, pilot_hashed_secret: str):
+    async def add_pilot_credentials(
+        self, pilot_id: int, pilot_hashed_secret: str
+    ) -> datetime:
 
         stmt = insert(PilotRegistrations).values(
             pilot_id=pilot_id, pilot_hashed_secret=pilot_hashed_secret
@@ -116,6 +125,22 @@ async def add_pilot_credentials(self, pilot_id: int, pilot_hashed_secret: str):
                     pilot_id=pilot_id, detail="this pilot has already credentials"
                 ) from e
 
+        added_creds = await self.get_pilot_creds_by_id(pilot_id)
+
+        return added_creds["PilotSecretCreationDate"]
+
+    async def set_pilot_credentials_expiration(
+        self, pilot_id: int, pilot_secret_expiration_date: DateTime
+    ):
+        #  Prepare the update statement
+        stmt = (
+            update(PilotRegistrations)
+            .values(pilot_secret_expiration_date=pilot_secret_expiration_date)
+            .where(PilotRegistrations.pilot_id == pilot_id)
+        )
+
+        await self.conn.execute(stmt)
+
     async def fetch_all_pilots(self):
         stmt = select(PilotRegistrations).with_for_update()
         result = await self.conn.execute(stmt)
@@ -134,3 +159,12 @@ async def get_pilot_by_reference(self, pilot_ref: str):
 
         # We assume it is unique...
         return dict((await self.conn.execute(stmt)).one()._mapping)
+
+    async def get_pilot_creds_by_id(self, pilot_id: int):
+        stmt = (
+            select(PilotRegistrations)
+            .with_for_update()
+            .where(PilotRegistrations.pilot_id == pilot_id)
+        )
+
+        return dict((await self.conn.execute(stmt)).one()._mapping)

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

@@ -3,7 +3,7 @@
 from sqlalchemy import DateTime, Double, ForeignKey, Index, Integer, String, Text
 from sqlalchemy.orm import declarative_base
 
-from ..utils import Column, EnumBackedBool, NullColumn
+from ..utils import Column, DateNowColumn, EnumBackedBool, NullColumn
 
 PilotAgentsDBBase = declarative_base()
 
@@ -64,3 +64,7 @@ class PilotRegistrations(PilotAgentsDBBase):
     )
     pilot_hashed_secret = Column("PilotHashedSecret", String(64))
     pilot_secret_use_count = Column("PilotSecretUseCount", Integer, default=0)
+    pilot_secret_creation_time = DateNowColumn("PilotSecretCreationDate")
+    pilot_secret_expiration_date = NullColumn(
+        "PilotSecretExpirationDate", DateTime(timezone=True)
+    )

diracx-db/tests/pilot_agents/test_pilot_agents_db.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from datetime import timedelta
+from time import sleep
+
 import pytest
 from sqlalchemy.exc import NoResultFound
 
@@ -74,10 +77,16 @@ async def test_create_pilot_and_verify_secret(pilot_agents_db: PilotAgentsDB):
         pilot_hashed_secret = hash(secret)
 
         # Add creds
-        await pilot_agents_db.add_pilot_credentials(
+        date_added = await pilot_agents_db.add_pilot_credentials(
             pilot_id=pilot_id, pilot_hashed_secret=pilot_hashed_secret
         )
 
+        expiration_date = date_added + timedelta(seconds=10)
+
+        await pilot_agents_db.set_pilot_credentials_expiration(
+            pilot_id=pilot_id, pilot_secret_expiration_date=expiration_date
+        )
+
         assert secret is not None
 
         await pilot_agents_db.verify_pilot_secret(
@@ -94,3 +103,46 @@ async def test_create_pilot_and_verify_secret(pilot_agents_db: PilotAgentsDB):
                 pilot_job_reference="I am a spider",
                 pilot_hashed_secret=pilot_hashed_secret,
             )
+
+
+async def test_create_pilot_and_verify_secret_with_delay(
+    pilot_agents_db: PilotAgentsDB,
+):
+
+    async with pilot_agents_db as pilot_agents_db:
+        pilot_reference = "pilot-reference-test"
+        # Register a pilot
+        await pilot_agents_db.add_pilot_references(
+            vo="lhcb",
+            pilot_ref=[pilot_reference],
+            grid_type="grid-type",
+        )
+
+        pilot = await pilot_agents_db.get_pilot_by_reference(pilot_reference)
+
+        pilot_id = pilot["PilotID"]
+
+        secret = "AW0nd3rfulS3cr3t"
+        pilot_hashed_secret = hash(secret)
+
+        # Add creds
+        date_added = await pilot_agents_db.add_pilot_credentials(
+            pilot_id=pilot_id, pilot_hashed_secret=pilot_hashed_secret
+        )
+
+        expiration_date = date_added + timedelta(seconds=1)
+
+        await pilot_agents_db.set_pilot_credentials_expiration(
+            pilot_id=pilot_id, pilot_secret_expiration_date=expiration_date
+        )
+
+        assert secret is not None
+
+        # So that the secret expires
+        sleep(3)
+
+        with pytest.raises(AuthorizationError):
+            await pilot_agents_db.verify_pilot_secret(
+                pilot_job_reference=pilot_reference,
+                pilot_hashed_secret=pilot_hashed_secret,
+            )

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

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+from datetime import timedelta
 from secrets import token_hex
 
+from diracx.core.settings import AuthSettings
 from diracx.db.sql import PilotAgentsDB
 
 # TODO: Move this hash function in diracx-logic, and rename it
@@ -13,23 +15,32 @@ def generate_pilot_secret() -> str:
     return token_hex(32)
 
 
-async def add_pilot_credentials(pilot_id: int, pilot_db: PilotAgentsDB) -> str:
+async def add_pilot_credentials(
+    pilot_id: int, pilot_db: PilotAgentsDB, settings: AuthSettings
+) -> str:
 
     # Get a random string
     # Can be customized
     random_secret = generate_pilot_secret()
 
     hashed_secret = hash(random_secret)
 
-    await pilot_db.add_pilot_credentials(
+    date_added = await pilot_db.add_pilot_credentials(
         pilot_id=pilot_id, pilot_hashed_secret=hashed_secret
     )
 
+    # Helps compatibility between sql engines
+    await pilot_db.set_pilot_credentials_expiration(
+        pilot_id=pilot_id,
+        pilot_secret_expiration_date=date_added  # type: ignore
+        + timedelta(seconds=settings.pilot_secret_expire_seconds),
+    )
+
     return random_secret
 
 
 def generate_pilot_scope(pilot: dict) -> str:
-    return f"vo:{pilot['VO']}"
+    return f"vo:{pilot['VO']} property:LimitedDelegation property:GenericPilot"
 
 
 async def try_login(

diracx-routers/tests/auth/test_pilot_auth.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from datetime import timedelta
+
 import pytest
 
 from diracx.db.sql.pilot_agents.db import PilotAgentsDB
@@ -54,10 +56,16 @@ async def test_create_pilot_and_verify_secret(test_client):
         pilot_id = pilot["PilotID"]
 
         # Add credentials to this pilot
-        await pilot_agents_db.add_pilot_credentials(
+        date_added = await pilot_agents_db.add_pilot_credentials(
             pilot_id=pilot_id, pilot_hashed_secret=pilot_hashed_secret
         )
 
+        expiration_date = date_added + timedelta(seconds=2)
+
+        await pilot_agents_db.set_pilot_credentials_expiration(
+            pilot_id=pilot_id, pilot_secret_expiration_date=expiration_date
+        )
+
     request_data = {"pilot_job_reference": pilot_reference, "pilot_secret": secret}
 
     r = test_client.post(
@@ -66,7 +74,7 @@ async def test_create_pilot_and_verify_secret(test_client):
         headers={"Content-Type": "application/json"},
     )
 
-    assert r.status_code == 200
+    assert r.status_code == 200, r.json()
 
     access_token = r.json()["access_token"]
     refresh_token = r.json()["refresh_token"]
@@ -109,7 +117,7 @@ async def test_create_pilot_and_verify_secret(test_client):
     )
 
     assert r.status_code == 401, r.json()
-    assert r.json()["detail"] == "bad pilot_id / pilot_secret"
+    assert r.json()["detail"] == "bad pilot_id / pilot_secret or secret has expired"
 
     # -----------------  Wrong ID  -----------------
     request_data = {"pilot_job_reference": "It is a reference", "pilot_secret": secret}