auth.py

from __future__ import annotations

import secrets
import time
from logging import getLogger
from uuid import uuid4

import aiohttp
import requests
from fastapi import APIRouter, Depends, HTTPException, Request, 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, Any

from murfey.server.murfey_db import murfey_db, url
from murfey.util.api import url_path_for
from murfey.util.config import get_security_config
from murfey.util.db import MurfeyUser as User, Session as MurfeySession

# Set up logger
logger = getLogger("murfey.server.api.auth")

# Set up router
router = APIRouter(
    prefix="/auth",
    tags=["Authentication"],
)


# Set up variables used for authentication
security_config = get_security_config()
auth_url = security_config.auth_url
ALGORITHM = security_config.auth_algorithm or "HS256"
SECRET_KEY = security_config.auth_key or secrets.token_hex(32)
oauth2_scheme = (
    OAuth2PasswordBearer(tokenUrl="auth/token")
    if security_config.auth_type == "password"
    else APIKeyCookie(name=security_config.cookie_key)
)
instrument_oauth2_scheme = (
    OAuth2PasswordBearer(tokenUrl="auth/token")
    if security_config.instrument_auth_type == "token"
    else lambda *args, **kwargs: None
)
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

instrument_server_tokens: dict[float, dict] = {}

# Set up database engine
try:
    _url = url(security_config)
    engine = create_engine(_url)
except Exception:
    engine = None


def hash_password(password: str) -> str:
    return pwd_context.hash(password)


"""
=======================================================================================
VALIDATION FUNCTIONS
=======================================================================================

Functions and helpers used to validate incoming requests from both the client and
the frontend.

'validate_token()' and 'validate_instrument_token()' are imported in the other FastAPI
modules and attached as dependencies to the routers. They validate the tokens passed
around internally by Murfey to ensure that the request is valid.

'validate_instrument_server_session_access()' and 'validate_frontend_session_access()'
are used to verify the IDs of sessions ot be accessed, and are attached as dependencies
to them.

'validate_user_instrument_access()' is used to verify the instrument server name being
accessed by the frontend, and is attached as a dependency as well.
"""

# Essential headers used for authentication to forward along if present
AUTH_HEADERS = (
    "authorization",
    "x-auth-request-access-token",
)


def check_user(username: str) -> bool:
    try:
        with Session(engine) as murfey_db:
            users = murfey_db.exec(select(User)).all()
    except Exception:
        return False
    return username in [u.username for u in users]


async def submit_to_auth_endpoint(
    url_subpath: str,
    request: Request,
    token: str,
) -> dict[str, Any]:
    """
    Helper function to forward incoming requests to an authentication server
    to verify that they are allowed to inspect the
    """

    # Forward only essentials auth-related headers
    headers = {
        key: value
        for key, value in dict(request.headers).items()
        if key.lower() in AUTH_HEADERS
    }
    if security_config.auth_type == "password":
        headers["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}/{url_subpath}",
            headers=headers,
        ) as response:
            success = response.status == 200
            validation_outcome: dict[str, Any] = await response.json()
    return validation_outcome if success and validation_outcome else {"valid": False}


async def validate_token(
    token: Annotated[str, Depends(oauth2_scheme)],
    request: Request,
):
    """
    Used by the backend routers to validate requests coming in from frontend.
    """
    try:
        # Validate using auth URL if provided; will error if invalid
        if auth_url:
            if not (
                await submit_to_auth_endpoint("validate_token", request, token)
            ).get("valid"):
                raise JWTError
        # 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])
            # 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,
            detail="Could not validate credentials from frontend",
            headers={"WWW-Authenticate": "Bearer"},
        )
    return None


def validate_session_against_visit(session_id: int, visit: str):
    """
    Checks that the session ID is associated with the claimed visit.
    """
    with Session(engine) as murfey_db:
        session_data = murfey_db.exec(
            select(MurfeySession).where(MurfeySession.id == session_id)
        ).all()
    if len(session_data) != 1:
        return False
    return visit == session_data[0].visit


async def validate_instrument_token(
    token: Annotated[str, Depends(instrument_oauth2_scheme)],
):
    """
    Used by the backend routers to check the incoming instrument server token.
    """
    try:
        # Validate using auth URL if provided
        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_token",
                    headers=headers,
                ) as response:
                    success = response.status == 200
                    validation_outcome = await response.json()
            if not (success and validation_outcome.get("valid")):
                raise JWTError
        else:
            # First, check if the token has expired
            decoded_data = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
            if expiry_time := decoded_data.get("expiry_time"):
                if expiry_time < time.time():
                    raise JWTError
            # Check that the decoded session corresponds to the visit
            elif decoded_data.get("session") is not None:
                if not validate_session_against_visit(
                    decoded_data["session"], decoded_data["visit"]
                ):
                    raise JWTError
            # Verify 'user' token if enabled
            elif security_config.allow_user_token:
                if not decoded_data.get("user"):
                    raise JWTError
            else:
                raise JWTError
    except JWTError:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Could not validate credentials from instrument",
            headers={"WWW-Authenticate": "Bearer"},
        )
    return None


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_instrument_server_session_access(
    session_id: int,
    token: Annotated[str, Depends(instrument_oauth2_scheme)],
) -> int:
    """
    Validates whether an instrument request can access information about this session
    """
    visit_name = get_visit_name(session_id)

    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


async def validate_frontend_session_access(
    session_id: int,
    request: Request,
    token: Annotated[str, Depends(oauth2_scheme)],
) -> int:
    """
    Validates whether a frontend request can access information about this session
    """
    visit_name = get_visit_name(session_id)
    if auth_url:
        if not (
            await submit_to_auth_endpoint(
                f"validate_visit_access/{visit_name}",
                request,
                token,
            )
        ).get("valid"):
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="You do not have access to this visit",
                headers={"WWW-Authenticate": "Bearer"},
            )
    return session_id


async def validate_user_instrument_access(
    instrument_name: str,
    request: Request,
    token: Annotated[str, Depends(oauth2_scheme)],
) -> str:
    """
    Validates whether a frontend request can access information about this instrument
    """
    if auth_url:
        if not (
            await submit_to_auth_endpoint(
                f"validate_instrument_access/{instrument_name}",
                request,
                token,
            )
        ).get("valid"):
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="You do not have access to this instrument",
                headers={"WWW-Authenticate": "Bearer"},
            )
    return instrument_name


# Create annotated session ID and instrument name for endpoints that need to verify them
MurfeySessionIDInstrument = Annotated[
    int, Depends(validate_instrument_server_session_access)
]
MurfeySessionIDFrontend = Annotated[int, Depends(validate_frontend_session_access)]
MurfeyInstrumentNameFrontend = Annotated[str, Depends(validate_user_instrument_access)]


"""
=======================================================================================
API ENDPOINTS AND HELPER FUNCTIONS/CLASSES
=======================================================================================
"""


def verify_password(plain_password: str, hashed_password: str) -> bool:
    return pwd_context.verify(plain_password, hashed_password)


def validate_user(username: str, password: str) -> bool:
    try:
        with Session(engine) as murfey_db:
            user = murfey_db.exec(select(User).where(User.username == username)).one()
    except Exception:
        return False
    return verify_password(password, user.hashed_password)


def create_access_token(data: dict, token: str = "") -> str:
    # 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}"},
            )
            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


class Token(BaseModel):
    access_token: str
    token_type: str


@router.post("/token")
async def generate_token(
    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
) -> Token:
    # 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)
            access_token = create_access_token(
                data={"user": form_data.username},
            )
        if not validated:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Incorrect username or password",
                headers={"WWW-Authenticate": "Bearer"},
            )
        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")
async def mint_session_token(session_id: MurfeySessionIDFrontend, db=murfey_db):
    visit = (
        db.exec(select(MurfeySession).where(MurfeySession.id == session_id)).one().visit
    )
    expiry_time = None
    if security_config.session_token_timeout:
        expiry_time = time.time() + security_config.session_token_timeout
    token = create_access_token(
        {
            "session": session_id,
            "visit": visit,
            "uuid": str(uuid4()),
            "expiry_time": expiry_time,
        }
    )
    return Token(access_token=token, token_type="bearer")


@router.get("/validate_token")
async def simple_token_validation(
    token: Annotated[str, Depends(validate_instrument_token)],
):
    return {"valid": True}