Simplify session token usage

Author: dormant-user

fastapiauthenticator/service.py

@@ -59,10 +59,6 @@ def __init__(
         elif isinstance(params, models.Parameters):
             self.params = [params]
 
-        self.route_map: Dict[str, models.Parameters] = {
-            param.path: param for param in self.params if param.route is APIRoute
-        }
-
         models.fallback.path = fallback_path
         models.fallback.button = fallback_button
 
@@ -95,19 +91,18 @@ def _verify_auth(
             Dict[str, str]:
             A dictionary containing the redirect URL to the secure path.
         """
-        utils.verify_login(
+        session_token = utils.verify_login(
             authorization=authorization,
             request=request,
             env_username=self.username,
             env_password=self.password,
         )
-        destination = request.cookies.get("X-Requested-By")
-        if parameter := self.route_map.get(destination):
+        if destination := request.cookies.get("X-Requested-By"):
             LOGGER.info("Setting session timeout for %s seconds", self.timeout)
             # Set session_token cookie with a timeout, to be used for session validation when redirected
             response.set_cookie(
                 key="session_token",
-                value=models.ws_session.client_auth[request.client.host].get("token"),
+                value=session_token,
                 httponly=True,
                 samesite="strict",
                 max_age=self.timeout,
@@ -118,7 +113,7 @@ def _verify_auth(
                 args=(request.client.host,),
                 interval=self.timeout,
             ).start()
-            return {"redirect_url": parameter.path}
+            return {"redirect_url": destination}
         raise HTTPException(
             status_code=status.HTTP_417_EXPECTATION_FAILED,
             detail="Unable to find secure route for the requested path.\n"
@@ -154,14 +149,14 @@ def _secure(self) -> None:
                 secure_route = APIWebSocketRoute(
                     path=param.path,
                     endpoint=param.function,
-                    dependencies=[Depends(utils.session_check)],
+                    dependencies=[Depends(utils.verify_session)],
                 )
             else:
                 secure_route = APIRoute(
                     path=param.path,
                     endpoint=param.function,
                     methods=["GET"],
-                    dependencies=[Depends(utils.session_check)],
+                    dependencies=[Depends(utils.verify_session)],
                 )
             self.app.routes.append(secure_route)
         self.app.routes.extend([login_route, session_route, verify_route, error_route])

fastapiauthenticator/utils.py

@@ -1,6 +1,6 @@
 import logging
 import secrets
-from typing import Dict, List, NoReturn, Union
+from typing import List, NoReturn
 
 from fastapi import status
 from fastapi.exceptions import HTTPException
@@ -73,17 +73,12 @@ def raise_error(request: Request) -> NoReturn:
     )
 
 
-def extract_credentials(
-    authorization: HTTPAuthorizationCredentials, host: str
-) -> List[str]:
+def extract_credentials(authorization: HTTPAuthorizationCredentials) -> List[str]:
     """Extract the credentials from ``Authorization`` headers and decode it before returning as a list of strings.
 
     Args:
         authorization: Authorization header from the request.
-        host: Host header from the request.
     """
-    if not authorization:
-        raise_error(host)
     decoded_auth = secure.base64_decode(authorization.credentials)
     # convert hex to a string
     auth = secure.hex_decode(decoded_auth)
@@ -95,7 +90,7 @@ def verify_login(
     request: Request,
     env_username: str,
     env_password: str,
-) -> Dict[str, Union[str, int]]:
+) -> str | NoReturn:
     """Verifies authentication and generates session token for each user.
 
     Args:
@@ -105,12 +100,13 @@ def verify_login(
         env_password: Environment variable for the password.
 
     Returns:
-        Dict[str, str]:
-        Returns a dictionary with the payload required to create the session token.
+        str:
+        Returns the session token.
     """
-    username, signature, timestamp = extract_credentials(
-        authorization, request.client.host
-    )
+    if authorization:
+        username, signature, timestamp = extract_credentials(authorization)
+    else:
+        raise_error(request)
     if secrets.compare_digest(username, env_username):
         hex_user = secure.hex_encode(env_username)
         hex_pass = secure.hex_encode(env_password)
@@ -122,15 +118,14 @@ def verify_login(
     if secrets.compare_digest(signature, expected_signature):
         models.ws_session.invalid[request.client.host] = 0
         key = secrets.token_urlsafe(64)
-        # fixme: By setting a path instead of timestamp, this can handle path specific sessions
-        models.ws_session.client_auth[request.client.host] = dict(
-            username=username, token=key, timestamp=int(timestamp)
-        )
-        return models.ws_session.client_auth[request.client.host]
+        models.ws_session.client_auth[request.client.host] = key
+        return key
     raise_error(request)
 
 
-def session_check(api_request: Request = None, api_websocket: WebSocket = None) -> None:
+def verify_session(
+    api_request: Request = None, api_websocket: WebSocket = None
+) -> None:
     """Check if the session is still valid.
 
     Args:
@@ -150,9 +145,7 @@ def session_check(api_request: Request = None, api_websocket: WebSocket = None)
             detail="Request or WebSocket connection is required for session check.",
         )
     session_token = request.cookies.get("session_token")
-    stored_token = models.ws_session.client_auth.get(request.client.host, {}).get(
-        "token"
-    )
+    stored_token = models.ws_session.client_auth.get(request.client.host)
     if (
         stored_token
         and session_token

verify/fast_api.py

@@ -1,12 +1,17 @@
 import uvicorn
 from fastapi import FastAPI, status
 from fastapi.requests import Request
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
 from fastapi.routing import APIRoute
 
 import fastapiauthenticator as auth
 
 
+def root_page() -> RedirectResponse:
+    """Re-direct the user to login page."""
+    return RedirectResponse(url="/sensitive-data", status_code=status.HTTP_302_FOUND)
+
+
 def hello_world() -> JSONResponse:
     """A simple function that returns a JSON response with a greeting message."""
     return JSONResponse({"message": "Hello, World!"})
@@ -22,6 +27,10 @@ def secure_function(_: Request) -> HTMLResponse:
 
 app = FastAPI(
     routes=[
+        APIRoute(
+            path="/",
+            endpoint=root_page,
+        ),
         APIRoute(
             path="/hello",
             endpoint=hello_world,