Add function to update the JSON Web Keys in the session

When new JSON Web Keys (JWKs) are fetched from the server, the
function added in this commit will help with updating them in the
session for subsequent uses, validating the tokens when needed.

1 files changed, 14 insertions, 0 deletions

diff --git a/uploader/session.py b/uploader/session.py
index 8b72bce..019d959 100644
--- a/uploader/session.py
+++ b/uploader/session.py
@@ -1,7 +1,9 @@
 """Deal with user sessions"""
 from uuid import UUID, uuid4
+from datetime import datetime
 from typing import Any, Optional, TypedDict
 
+from authlib.jose import KeySet
 from flask import request, session
 from pymonad.either import Left, Right, Either
 
@@ -89,3 +91,15 @@ def user_details() -> UserDetails:
 def user_token() -> Either:
     """Retrieve the user token."""
     return session_info()["user"]["token"]
+
+
+def set_auth_server_jwks(keyset: KeySet) -> KeySet:
+    """Update the JSON Web Keys in the session."""
+    save_session_info({
+        **session_info(),
+        "auth_server_jwks": {
+            "last-updated": datetime.now().timestamp(),
+            "jwks": keyset.as_dict()
+        }
+    })
+    return keyset