Include the ability to handle multiple sessions

Author: dormant-user

fastapiauthenticator/__init__.py

@@ -1,5 +1,5 @@
 from fastapiauthenticator.enums import APIEndpoints, APIMethods  # noqa: F401,E402
-from fastapiauthenticator.models import Params  # noqa: F401,E402
+from fastapiauthenticator.models import Parameters  # noqa: F401,E402
 from fastapiauthenticator.service import Authenticator  # noqa: F401,E402
 from fastapiauthenticator.version import version  # noqa: F401,E402
 

fastapiauthenticator/models.py

@@ -3,34 +3,31 @@
 
 from fastapi.routing import APIRoute, APIWebSocketRoute
 from fastapi.templating import Jinja2Templates
-from pydantic import BaseModel, Field, PositiveInt
+from pydantic import BaseModel, Field
 
 from fastapiauthenticator.enums import APIMethods
 
 templates = Jinja2Templates(directory=pathlib.Path(__file__).parent / "templates")
 
 
-class Params(BaseModel):
+class Parameters(BaseModel):
     """Parameters for the Authenticator class.
 
-    >>> Params
+    >>> Parameters
 
     Attributes:
+        path: Path for the secure route, must start with '/'.
         function: Function to be called for secure routes after authentication.
         methods: List of HTTP methods that the secure function will handle.
         route: Type of route to be used for secure routes, either APIWebSocketRoute or APIRoute.
-        path: Path for the secure route, must start with '/'.
     """
 
-    function: Callable
-    methods: List[APIMethods] = None
-    route: Type[APIWebSocketRoute] | Type[APIRoute]
     path: str = Field(
         pattern="^/.*$", description="Path for the secure route, must start with '/'"
     )
-    timeout: PositiveInt = Field(
-        ge=0, default=300, description="Session timeout in seconds."
-    )
+    function: Callable
+    methods: List[APIMethods] = [APIMethods.GET]
+    route: Type[APIWebSocketRoute] | Type[APIRoute] = APIRoute
 
 
 class WSSession(BaseModel):
@@ -71,15 +68,17 @@ class RedirectException(Exception):
         https://fastapi.tiangolo.com/tutorial/handling-errors/#install-custom-exception-handlers
     """
 
-    def __init__(self, location: str, detail: Optional[str] = ""):
+    def __init__(self, source: str, destination: str, detail: Optional[str] = ""):
         """Instantiates the ``RedirectException`` object with the required parameters.
 
         Args:
-            location: Location for redirect.
+            source: Source from where the redirect is initiated.
+            destination: Location to redirect.
             detail: Reason for redirect.
         """
-        self.location = location
         self.detail = detail
+        self.source = source
+        self.destination = destination
 
 
 ws_session = WSSession()

fastapiauthenticator/service.py

@@ -1,10 +1,9 @@
 import logging
 import os
-from threading import Timer
 from typing import Dict, List
 
 import dotenv
-from fastapi import FastAPI
+from fastapi.applications import FastAPI
 from fastapi.params import Depends
 from fastapi.requests import Request
 from fastapi.responses import Response
@@ -29,7 +28,8 @@ class Authenticator:
     def __init__(
         self,
         app: FastAPI,
-        params: models.Params | List[models.Params],
+        params: models.Parameters | List[models.Parameters],
+        timeout: int = 300,
         username: str = os.environ.get("USERNAME"),
         password: str = os.environ.get("PASSWORD"),
         fallback_button: str = models.fallback.button,
@@ -39,6 +39,8 @@ def __init__(
 
         Args:
             app: FastAPI application instance to which the authenticator will be added.
+            params: Parameters for the secure routes, can be a single `Parameters` object or a list of `Parameters`.
+            timeout: Session timeout in seconds, default is 300 seconds (5 minutes).
             username: Username for authentication, can be set via environment variable 'USERNAME'.
             password: Password for authentication, can be set via environment variable 'PASSWORD'.
             fallback_button: Title for the fallback button, defaults to "LOGIN".
@@ -51,10 +53,10 @@ def __init__(
 
         if isinstance(params, list):
             self.params = params
-        elif isinstance(params, models.Params):
+        elif isinstance(params, models.Parameters):
             self.params = [params]
 
-        self.route_map: Dict[str, models.Params] = {
+        self.route_map: Dict[str, models.Parameters] = {
             param.path: param for param in self.params if param.route is APIRoute
         }
 
@@ -69,6 +71,7 @@ def __init__(
 
         self.username = username
         self.password = password
+        self.timeout = timeout
 
         self._secure()
 
@@ -95,84 +98,28 @@ def _verify_auth(
             env_username=self.username,
             env_password=self.password,
         )
-        referer = request.headers.get("Referer")
-        origin = request.headers.get("Origin")
-        destination = referer.replace(origin, "")
+        destination = request.cookies.get("X-Requested-By")
         parameter = self.route_map.get(destination)
-        private_route = APIRoute(
-            path=parameter.path,
-            endpoint=parameter.function,
-            methods=parameter.methods,
-            dependencies=[Depends(utils.session_check)],
-        )
-        for route in self.app.routes:
-            if route.path == private_route.path:
-                LOGGER.info(
-                    "Route %s already exists, removing it to replace with secure route.",
-                    private_route.path,
-                )
-                self.app.routes.remove(route)
-                break
-        self.app.routes.append(private_route)
-        LOGGER.info("Setting session timeout for %s seconds", parameter.timeout)
-        self._handle_session(
-            response=response,
-            request=request,
-            secure_route=private_route,
-            timeout=parameter.timeout,
-        )
-        return {"redirect_url": parameter.path}
-
-    def _setup_session_route(self, secure_route: APIRoute) -> None:
-        """Removes the secure route and adds a routing logic for invalid sessions.
-
-        Args:
-            secure_route: Secure route to be removed from the app after the session timeout.
-        """
-        LOGGER.info("Session expired, removing secure route: %s", secure_route.path)
-        self.app.routes.remove(secure_route)
-        LOGGER.info(
-            "Adding session route to handle expired sessions at %s", secure_route.path
-        )
-        self.app.routes.append(
-            APIRoute(
-                path=secure_route.path,
-                endpoint=endpoints.session,
-                methods=["GET"],
-            )
-        )
-
-    def _handle_session(
-        self,
-        response: Response,
-        request: Request,
-        secure_route: APIRoute,
-        timeout: int,
-    ) -> None:
-        """Handle session management by setting a cookie and scheduling session removal.
-
-        Args:
-            response: Response object to set the session cookie.
-            request: Request object containing client information.
-            secure_route: Secure route to be removed from the app after the session timeout.
-        """
-        # Remove the secure route after the session timeout - backend
-        Timer(
-            function=self._setup_session_route,
-            args=(secure_route,),
-            interval=timeout,
-        ).start()
-        # Set the max age in session cookie to session timeout - frontend
+        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"),
             httponly=True,
             samesite="strict",
-            max_age=timeout,
+            max_age=self.timeout,
         )
+        # todo: Session should be cleared at client side after timeout
+        response.delete_cookie(key="X-Requested-By")
+        return {"redirect_url": parameter.path}
 
     def _secure(self) -> None:
         """Create the login and verification routes for the APIAuthenticator."""
+        login_route = APIRoute(
+            path=enums.APIEndpoints.fastapi_login,
+            endpoint=endpoints.login,
+            methods=["GET"],
+        )
         error_route = APIRoute(
             path=enums.APIEndpoints.fastapi_error,
             endpoint=endpoints.error,
@@ -190,6 +137,7 @@ def _secure(self) -> None:
         )
         for param in self.params:
             if param.route is APIWebSocketRoute:
+                # WebSocket routes will not have a login path, they will be protected by session check
                 secure_route = APIWebSocketRoute(
                     path=param.path,
                     endpoint=param.function,
@@ -198,8 +146,9 @@ def _secure(self) -> None:
             else:
                 secure_route = APIRoute(
                     path=param.path,
-                    endpoint=endpoints.login,
+                    endpoint=param.function,
                     methods=["GET"],
+                    dependencies=[Depends(utils.session_check)],
                 )
             self.app.routes.append(secure_route)
-        self.app.routes.extend([session_route, verify_route, error_route])
+        self.app.routes.extend([login_route, session_route, verify_route, error_route])

fastapiauthenticator/utils.py

@@ -2,11 +2,12 @@
 import secrets
 from typing import Dict, List, NoReturn, Union
 
-from fastapi import WebSocket, status
+from fastapi import status
 from fastapi.exceptions import HTTPException
 from fastapi.requests import Request
 from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
 from fastapi.security import HTTPAuthorizationCredentials
+from fastapi.websockets import WebSocket
 
 from fastapiauthenticator import enums, models, secure
 
@@ -23,8 +24,9 @@ def failed_auth_counter(host: str) -> None:
         models.ws_session.invalid[host] += 1
     except KeyError:
         models.ws_session.invalid[host] = 1
-    if models.ws_session.invalid[host] >= 3:
-        raise models.RedirectException(location=enums.APIEndpoints.fastapi_error)
+    # todo: fix this
+    # if models.ws_session.invalid[host] >= 3:
+    #     raise models.RedirectException(location=enums.APIEndpoints.fastapi_error)
 
 
 def redirect_exception_handler(
@@ -44,14 +46,17 @@ def redirect_exception_handler(
     LOGGER.warning("Exception cookies: %s", request.cookies)
     if request.url.path == enums.APIEndpoints.fastapi_verify_login:
         response = JSONResponse(
-            content={"redirect_url": exception.location}, status_code=200
+            content={"redirect_url": exception.destination}, status_code=200
         )
     else:
-        response = RedirectResponse(url=exception.location)
+        response = RedirectResponse(url=exception.destination)
     if exception.detail:
         response.set_cookie(
             "detail", exception.detail.upper(), httponly=True, samesite="strict"
         )
+    response.set_cookie(
+        "X-Requested-By", exception.source, httponly=True, samesite="strict"
+    )
     return response
 
 
@@ -121,43 +126,40 @@ def verify_login(
     raise_error(host)
 
 
-def session_check(request: Request = None, websocket: WebSocket = None) -> None:
+def session_check(api_request: Request = None, api_websocket: WebSocket = None) -> None:
     """Check if the session is still valid.
 
     Args:
-        request: Request containing client information.
-        websocket: WebSocket connection object.
+        api_request: Request containing client information.
+        api_websocket: WebSocket connection object.
 
     Raises:
         HTTPException: If the session is invalid or expired.
     """
-    if request:
-        host = request.client.host
-        session_token = request.cookies.get("session_token")
-    elif websocket:
-        host = websocket.client.host
-        session_token = websocket.cookies.get("session_token")
+    if api_request:
+        request = api_request
+    elif api_websocket:
+        request = api_websocket
     else:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail="Request or WebSocket connection is required for session check.",
         )
-    stored_token = models.ws_session.client_auth.get(host, {}).get("token")
+    session_token = request.cookies.get("session_token")
+    stored_token = models.ws_session.client_auth.get(request.client.host, {}).get(
+        "token"
+    )
     if (
         stored_token
         and session_token
         and secrets.compare_digest(session_token, stored_token)
     ):
-        LOGGER.info("Session is valid for host: %s", host)
+        LOGGER.info("Session is valid for host: %s", request.client.host)
         return
-    # todo: this will fail all new sessions
-    #   the auth page route will be removed from the app if session1 is valid
-    #   when session2 is created, the auth page route will not be available, and since session2 is not authenticated,
-    #   the content cannot be rendered
-    LOGGER.warning("Session is invalid for host: %s", host)
+    LOGGER.warning("Session is invalid or expired for host: %s", request.client.host)
     raise models.RedirectException(
-        location=enums.APIEndpoints.fastapi_session,
-        detail="Session expired or invalid. Please log in again.",
+        source=request.url.path,
+        destination=enums.APIEndpoints.fastapi_login,
     )
 
 
@@ -172,10 +174,6 @@ def clear_session(request: Request, response: HTMLResponse) -> HTMLResponse:
         HTMLResponse:
         Returns the response object with the session token cleared.
     """
-    for cookie in request.cookies:
-        # Deletes all cookies stored in current session
-        LOGGER.info("Deleting cookie: '%s'", cookie)
-        response.delete_cookie(cookie)
     response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
     response.headers["Authorization"] = ""
     return response

verify/index.html

@@ -45,7 +45,7 @@
       amplitude: 0,
       frequency: 6,
       color: "#ffffff",
-      autostart: true,
+      autostart: false,
       cover: true,
     });