From d3b1a05299b1c0675a66fbbbfec689091a11963e Mon Sep 17 00:00:00 2001
From: dr-frmr <docterformer@protonmail.com>
Date: Thu, 3 Apr 2025 16:24:23 -0400
Subject: [PATCH] add JWT expiry

---
 src/routers/keys.py       |  3 +++
 src/security.py           | 10 ++++++++++
 tests/routes/test_keys.py | 19 +++++++++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/src/routers/keys.py b/src/routers/keys.py
index aaadd82e..c6ee07bd 100644
--- a/src/routers/keys.py
+++ b/src/routers/keys.py
@@ -1,3 +1,4 @@
+import datetime
 import logging
 import os
 
@@ -23,6 +24,7 @@ router = APIRouter(
 
 @router.post("")
 async def create_key(
+    expires_at: datetime.datetime | None = None,
     app_id: str | None = None,
     user_id: str | None = None,
     session_id: str | None = None,
@@ -40,6 +42,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 f0aae28c..78549fdb 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