Merge pull request #21 from UpstreamDataInc/dev_session

Improve session handling

Author: b-rowan

goosebit/auth/__init__.py

@@ -4,7 +4,8 @@
 from fastapi.requests import Request
 from fastapi.security import SecurityScopes
 from fastapi.websockets import WebSocket
-from jose import jwt
+from joserfc import jwt
+from joserfc.errors import BadSignatureError
 
 from goosebit.settings import PWD_CXT, SECRET, USERS
 
@@ -29,46 +30,49 @@ async def authenticate_user(request: Request):
     return user
 
 
-sessions = {}
-
-
-def create_session(email: str) -> str:
-    token = jwt.encode({"email": email}, SECRET)
-    sessions[token] = email
-    return token
+def create_session(username: str) -> str:
+    return jwt.encode(
+        header={"alg": "HS256"}, claims={"username": username}, key=SECRET
+    )
 
 
 def authenticate_session(request: Request):
     session_id = request.cookies.get("session_id")
-    if session_id is None or session_id not in sessions:
+    if session_id is None:
         raise HTTPException(
             status_code=302,
             headers={"location": str(request.url_for("login"))},
             detail="Invalid session ID",
         )
     user = get_user_from_session(session_id)
+    if user is None:
+        raise HTTPException(
+            status_code=302,
+            headers={"location": str(request.url_for("login"))},
+            detail="Invalid username",
+        )
     return user
 
 
 def authenticate_api_session(request: Request):
     session_id = request.cookies.get("session_id")
-    if session_id is None or session_id not in sessions:
-        raise HTTPException(status_code=401, detail="Not logged in")
     user = get_user_from_session(session_id)
+    if user is None:
+        raise HTTPException(status_code=401, detail="Not logged in")
     return user
 
 
 def authenticate_ws_session(websocket: WebSocket):
     session_id = websocket.cookies.get("session_id")
-    if session_id is None or session_id not in sessions:
-        raise HTTPException(status_code=401, detail="Not logged in")
     user = get_user_from_session(session_id)
+    if user is None:
+        raise HTTPException(status_code=401, detail="Not logged in")
     return user
 
 
 def auto_redirect(request: Request):
     session_id = request.cookies.get("session_id")
-    if session_id is None or session_id not in sessions:
+    if get_user_from_session(session_id) is None:
         return request
     raise HTTPException(
         status_code=302,
@@ -78,16 +82,20 @@ def auto_redirect(request: Request):
 
 
 def get_user_from_session(session_id: str):
-    for username in USERS:
-        if username == sessions.get(session_id):
-            return username
+    if session_id is None:
+        return
+    try:
+        session_data = jwt.decode(session_id, SECRET)
+        return session_data.claims["username"]
+    except (BadSignatureError, LookupError):
+        pass
 
 
 def get_current_user(request: Request):
     session_id = request.cookies.get("session_id")
-    if session_id is None or session_id not in sessions:
-        return None
     user = get_user_from_session(session_id)
+    if user is None:
+        return None
     return USERS[user]
 
 

poetry.lock

@@ -11,12 +11,12 @@ fastapi = {extras = ["uvicorn"], version = "^0.111.0"}
 python-multipart = "^0.0.9"
 jinja2 = "^3.1.4"
 itsdangerous = "^2.2.0"
-python-jose = "^3.3.0"
 tortoise-orm = "^0.21.4"
 aerich = "^0.7.2"
 aiofiles = "^24.1.0"
 websockets = "^12.0"
 argon2-cffi = "^23.1.0"
+joserfc = "^1.0.0"
 semver = "^3.0.2"
 
 [tool.poetry.group.dev.dependencies]