diff --git a/src/routers/keys.py b/src/routers/keys.py index 87705559..c40b0ae7 100644 --- a/src/routers/keys.py +++ b/src/routers/keys.py @@ -1,3 +1,4 @@ +import datetime import logging import os @@ -31,6 +32,7 @@ async def create_key( collection_id: str | None = Query( None, description="ID of the collection to scope the key to" ), + expires_at: datetime.datetime | None = None, ): """Create a new Key""" if not USE_AUTH: @@ -44,6 +46,7 @@ async def create_key( key_str = create_jwt( JWTParams( + exp=expires_at.isoformat() if expires_at else None, ap=app_id, us=user_id, se=session_id, diff --git a/src/security.py b/src/security.py index 9b018473..6b1e837e 100644 --- a/src/security.py +++ b/src/security.py @@ -54,6 +54,7 @@ class JWTParams(BaseModel): Fields (all optional other than `t`): `t`: a string timestamp of when the JWT was created + `exp`: a string timestamp of when the JWT expires (optional) `ad`: a boolean flag indicating if the JWT is an admin JWT `ap`: (string) app id `us`: (string) user id @@ -62,6 +63,7 @@ class JWTParams(BaseModel): """ t: str = datetime.datetime.now().isoformat() + exp: Optional[str] = None ad: Optional[bool] = None ap: Optional[str] = None us: Optional[str] = None @@ -92,6 +94,14 @@ async def verify_jwt(token: str) -> JWTParams: ) if "t" in decoded: params.t = decoded["t"] + if "exp" in decoded: + params.exp = decoded["exp"] + if ( + params.exp + and datetime.datetime.fromisoformat(params.exp) + < datetime.datetime.now() + ): + raise AuthenticationException("JWT expired") if "ad" in decoded: params.ad = decoded["ad"] if "ap" in decoded: diff --git a/tests/routes/test_keys.py b/tests/routes/test_keys.py index 5e0ec01c..d7690e46 100644 --- a/tests/routes/test_keys.py +++ b/tests/routes/test_keys.py @@ -42,3 +42,22 @@ def test_create_key_with_params(auth_client, sample_data): ) assert response.status_code == 200 assert "key" in response.json() + + +def test_create_key_with_expires_at(auth_client, sample_data): + """Test creating a key with an expiration date""" + response = auth_client.post("/v1/keys", params={"expires_at": "2025-01-01"}) + + # Only admin JWT should be allowed + if auth_client.auth_type == "admin": + # key with no params should fail + assert response.status_code == 422 + return + else: + assert response.status_code == 401 + + test_app, _ = sample_data + + # assert that the key is expired + response = auth_client.post("/v1/keys", params={"app_id": test_app.public_id}) + assert response.status_code == 401