Used FastAPI's 'APIKeyCookie' object for cookie authentication instead of homebrew method; split session access validation function for instrument server and frontend into separate functions; instrument validation function was incorrectly calling oauth2 scheme for frontend instead of backend; fixed logic for 'create_access_token' and 'generate_token' for handling authentication using either 'password' or 'cookie'; 'simple_token_validation()' should be using instrument server validation function instead

Author: tieneupin

src/murfey/server/api/auth.py

@@ -3,17 +3,22 @@
 import secrets
 import time
 from logging import getLogger
-from typing import Annotated, Dict
+from typing import Dict
 from uuid import uuid4
 
 import aiohttp
 import requests
-from fastapi import APIRouter, Depends, HTTPException, Request, status
-from fastapi.security import HTTPBearer, OAuth2PasswordBearer, OAuth2PasswordRequestForm
+from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi.security import (
+    APIKeyCookie,
+    OAuth2PasswordBearer,
+    OAuth2PasswordRequestForm,
+)
 from jose import JWTError, jwt
 from passlib.context import CryptContext
 from pydantic import BaseModel
 from sqlmodel import Session, create_engine, select
+from typing_extensions import Annotated
 
 from murfey.server.murfey_db import murfey_db, url
 from murfey.util.api import url_path_for
@@ -31,38 +36,6 @@
 )
 
 
-class CookieScheme(HTTPBearer):
-    def __init__(
-        self,
-        *,
-        description: str | None = None,
-        auto_error: bool = True,
-        cookie_key: str = "cookie_auth",
-    ):
-        """
-        Args:
-            cookie_key: Cookie key to look for in requests
-        """
-        super().__init__(
-            description=description,
-            auto_error=auto_error,
-        )
-
-        self.cookie_key = cookie_key
-
-    async def __call__(self, request: Request):
-        token = request.cookies.get(self.cookie_key)
-        if token is None:
-            if self.auto_error:
-                raise HTTPException(
-                    status_code=status.HTTP_401_UNAUTHORIZED,
-                    detail="Not authenticated",
-                )
-            else:
-                return None
-        return token
-
-
 # Set up variables used for authentication
 security_config = get_security_config()
 auth_url = security_config.auth_url
@@ -71,7 +44,7 @@ async def __call__(self, request: Request):
 if security_config.auth_type == "password":
     oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 else:
-    oauth2_scheme = CookieScheme(cookie_key=security_config.cookie_key)
+    oauth2_scheme = APIKeyCookie(name=security_config.cookie_key)
 if security_config.instrument_auth_type == "token":
     instrument_oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 else:
@@ -138,19 +111,19 @@ async def validate_token(token: Annotated[str, Depends(oauth2_scheme)]):
                     validation_outcome = await response.json()
             if not (success and validation_outcome.get("valid")):
                 raise JWTError
-        # Auth URL MUST be provided if authenticating using cookies
+        # If authenticating using cookies; an auth URL MUST be provided
         else:
             if security_config.auth_type == "cookie":
                 raise JWTError
-
         # Validate using password
         if security_config.auth_type == "password":
             decoded_data = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
-            # Frontend validation
+            # Check that the user is present and is valid
             if decoded_data.get("user"):
                 if not check_user(decoded_data["user"]):
                     raise JWTError
-
+            else:
+                raise JWTError
     except JWTError:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
@@ -221,7 +194,7 @@ async def validate_instrument_token(
 
 """
 =======================================================================================
-VALIDATING SESSION IDS
+SESSION ID VALIDATION
 =======================================================================================
 
 Annotated ints are defined here that trigger validation of the session IDs in incoming
@@ -233,104 +206,87 @@ async def validate_instrument_token(
 """
 
 
-async def validate_visit(visit_name: str, token: str, instrument_access: bool) -> bool:
-    """
-    Validates the incoming token depending on whether it's from the instrument or the frontend
-    """
-
-    # If this token is from the instrument server, validate this way
-    if instrument_access:
-        if security_config.instrument_auth_url:
-            async with aiohttp.ClientSession() as session:
-                headers = (
-                    {}
-                    if not security_config.instrument_auth_type
-                    else {"Authorization": f"Bearer {token}"}
-                )
-                async with session.get(
-                    f"{security_config.instrument_auth_url}/validate_visit_access/{visit_name}",
-                    headers=headers,
-                ) as response:
-                    success = response.status == 200
-                    validation_outcome = await response.json()
-            if not (success and validation_outcome.get("valid")):
-                return False
-    # Otherwise, use this validation method
-    else:
-        if security_config.auth_url:
-            headers = (
-                {}
-                if security_config.auth_type == "cookie"
-                else {"Authorization": f"Bearer {token}"}
-            )
-            cookies = (
-                {security_config.cookie_key: token}
-                if security_config.auth_type == "cookie"
-                else {}
-            )
-            async with aiohttp.ClientSession(cookies=cookies) as session:
-                async with session.get(
-                    f"{auth_url}/validate_visit_access/{visit_name}",
-                    headers=headers,
-                ) as response:
-                    success = response.status == 200
-                    validation_outcome = await response.json()
-            if not (success and validation_outcome.get("valid")):
-                return False
-    return True
+def get_visit_name(session_id: int) -> str:
+    with Session(engine) as murfey_db:
+        return (
+            murfey_db.exec(select(MurfeySession).where(MurfeySession.id == session_id))
+            .one()
+            .visit
+        )
 
 
-async def validate_session_access(
+async def validate_frontend_session_access(
     session_id: int,
     token: Annotated[str, Depends(oauth2_scheme)],
-    instrument_access: bool,
 ) -> int:
     """
-    Validates whether the request is authorised to access information about this session
+    Validates whether a frontend request can access information about this session
     """
-    with Session(engine) as murfey_db:
-        visit_name = (
-            murfey_db.exec(select(MurfeySession).where(MurfeySession.id == session_id))
-            .one()
-            .visit
+    visit_name = get_visit_name(session_id)
+
+    if auth_url:
+        headers = (
+            {}
+            if security_config.auth_type == "cookie"
+            else {"Authorization": f"Bearer {token}"}
         )
-    validated = await validate_visit(visit_name, token, instrument_access)
-    if not validated:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="You do not have access to this visit",
-            headers={"WWW-Authenticate": "Bearer"},
+        cookies = (
+            {security_config.cookie_key: token}
+            if security_config.auth_type == "cookie"
+            else {}
         )
+        async with aiohttp.ClientSession(cookies=cookies) as session:
+            async with session.get(
+                f"{auth_url}/validate_visit_access/{visit_name}",
+                headers=headers,
+            ) as response:
+                success = response.status == 200
+                validation_outcome: dict = await response.json()
+        if not (success and validation_outcome.get("valid")):
+            logger.warning("Unauthorised visit access request from frontend")
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="You do not have access to this visit",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
     return session_id
 
 
-def validate_session_access_wrapper(instrument_access: bool):
+async def validate_instrument_session_access(
+    session_id: int,
+    token: Annotated[str, Depends(instrument_oauth2_scheme)],
+) -> int:
     """
-    Factory that returns an async wrapper arond 'validate_session_access' with the
-    'instrument_access' field preconfigured.
-
-    This is used to generate FastAPI-compatible dependencies for validating session
-    access, based on the context of the request.
-
-    Unlike 'functools.partial', this approach preserves introspection compatibility
-    required by FastAPI for dependency resolution and OpenAPI generation.
+    Validates whether an instrument request can access information about this session
     """
+    visit_name = get_visit_name(session_id)
 
-    async def _validate(session_id: int, token: Annotated[str, Depends(oauth2_scheme)]):
-        return await validate_session_access(
-            session_id, token, instrument_access=instrument_access
-        )
-
-    return _validate
+    if security_config.instrument_auth_url:
+        async with aiohttp.ClientSession() as session:
+            headers = (
+                {}
+                if not security_config.instrument_auth_type
+                else {"Authorization": f"Bearer {token}"}
+            )
+            async with session.get(
+                f"{security_config.instrument_auth_url}/validate_visit_access/{visit_name}",
+                headers=headers,
+            ) as response:
+                success = response.status == 200
+                validation_outcome = await response.json()
+        if not (success and validation_outcome.get("valid")):
+            logger.warning("Unauthorised visit access request from instrument")
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="You do not have access to this visit",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+    return session_id
 
 
 # Set validation conditions for the session ID based on where the request is from
-MurfeySessionIDFrontend = Annotated[
-    int, Depends(validate_session_access_wrapper(instrument_access=False))
-]
-MurfeySessionIDInstrument = Annotated[
-    int, Depends(validate_session_access_wrapper(instrument_access=True))
-]
+MurfeySessionIDFrontend = Annotated[int, Depends(validate_frontend_session_access)]
+MurfeySessionIDInstrument = Annotated[int, Depends(validate_instrument_session_access)]
 
 
 """
@@ -354,23 +310,27 @@ def validate_user(username: str, password: str) -> bool:
 
 
 def create_access_token(data: dict, token: str = "") -> str:
-    if auth_url and data.get("session"):
-        session_id = data["session"]
-        if not isinstance(session_id, int) and session_id > 0:
-            # check the session ID is alphanumeric for security
-            raise ValueError("Session ID was invalid (not alphanumeric)")
-        minted_token_response = requests.get(
-            f"{auth_url}{url_path_for('auth.router', 'mint_session_token', session_id=session_id)}",
-            headers={"Authorization": f"Bearer {token}"},
-        )
-        if minted_token_response.status_code != 200:
-            raise RuntimeError(
-                f"Request received status code {minted_token_response.status_code} when trying to create session token"
+
+    # If authenticating with password, auth URL needs a 'mint_session_token' endpoint
+    if security_config.auth_type == "password":
+        if auth_url and data.get("session"):
+            session_id = data["session"]
+            if not isinstance(session_id, int) and session_id > 0:
+                # Check the session ID is alphanumeric for security
+                raise ValueError("Session ID was invalid (not alphanumeric)")
+            minted_token_response = requests.get(
+                f"{auth_url}{url_path_for('auth.router', 'mint_session_token', session_id=session_id)}",
+                headers={"Authorization": f"Bearer {token}"},
             )
-        return minted_token_response.json()["access_token"]
+            if minted_token_response.status_code != 200:
+                raise RuntimeError(
+                    f"Request received status code {minted_token_response.status_code} when trying to create session token"
+                )
+            return minted_token_response.json()["access_token"]
 
     to_encode = data.copy()
 
+    # Make token for instrument
     encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
     return encoded_jwt
 
@@ -384,31 +344,36 @@ class Token(BaseModel):
 async def generate_token(
     form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
 ) -> Token:
-    if auth_url:
-        data = aiohttp.FormData()
-        data.add_field("username", form_data.username)
-        data.add_field("password", form_data.password)
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                f"{auth_url}{url_path_for('auth.router', 'generate_token')}",
-                data=data,
-            ) as response:
-                validated = response.status == 200
-                token = await response.json()
-                access_token = token.get("access_token")
-    else:
-        validated = validate_user(form_data.username, form_data.password)
-    if not validated:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Incorrect username or password",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    if not auth_url:
-        access_token = create_access_token(
-            data={"user": form_data.username},
-        )
-    return Token(access_token=access_token, token_type="bearer")
+    # Only generate a token if it's a password
+    if security_config.auth_type == "password":
+        if auth_url:
+            data = aiohttp.FormData()
+            data.add_field("username", form_data.username)
+            data.add_field("password", form_data.password)
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    f"{auth_url}{url_path_for('auth.router', 'generate_token')}",
+                    data=data,
+                ) as response:
+                    validated = response.status == 200
+                    token = await response.json()
+                    access_token = token.get("access_token")
+        else:
+            validated = validate_user(form_data.username, form_data.password)
+        if not validated:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Incorrect username or password",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        if not auth_url:
+            access_token = create_access_token(
+                data={"user": form_data.username},
+            )
+        return Token(access_token=access_token, token_type="bearer")
+
+    # Return empty token otherwise
+    return Token(access_token="", token_type="bearer")
 
 
 @router.get("/sessions/{session_id}/token")
@@ -431,5 +396,7 @@ async def mint_session_token(session_id: MurfeySessionIDFrontend, db=murfey_db):
 
 
 @router.get("/validate_token")
-async def simple_token_validation(token: Annotated[str, Depends(validate_token)]):
+async def simple_token_validation(
+    token: Annotated[str, Depends(validate_instrument_token)]
+):
     return {"valid": True}