homelable/backend/app/core/security.py

import hmac
import secrets
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from functools import lru_cache
from typing import Any

import bcrypt
import jwt
from authlib.integrations.starlette_client import OAuth, StarletteOAuth2App
from fastapi import Request, Response, status
from jwt import InvalidTokenError
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.responses import JSONResponse

from app.core.config import origin_of, settings

ACCESS_TOKEN_USE = "access"
OIDC_SESSION_TOKEN_USE = "oidc_session"
SAFE_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "TRACE"})


@dataclass(frozen=True)
class OIDCSession:
    subject: str
    issuer: str
    display_name: str
    csrf_token: str


def verify_password(plain: str, hashed: str) -> bool:
    if not plain or not hashed:
        return False
    try:
        return bcrypt.checkpw(plain.encode("utf-8"), hashed.encode("utf-8"))
    except (ValueError, TypeError):
        return False


def hash_password(password: str) -> str:
    return bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()).decode("utf-8")


def create_access_token(subject: str) -> str:
    expire = datetime.now(timezone.utc) + timedelta(minutes=settings.access_token_expire_minutes)
    payload = {"sub": subject, "exp": expire, "token_use": ACCESS_TOKEN_USE}
    return str(jwt.encode(payload, settings.secret_key, algorithm=settings.algorithm))


def decode_token(token: str) -> str | None:
    try:
        payload = jwt.decode(token, settings.secret_key, algorithms=[settings.algorithm])
        if payload.get("token_use") not in (None, ACCESS_TOKEN_USE):
            return None
        sub = payload.get("sub")
        return str(sub) if sub is not None else None
    except InvalidTokenError:
        return None


def create_oidc_session_token(userinfo: Mapping[str, Any]) -> str:
    subject = userinfo.get("sub")
    issuer = userinfo.get("iss")
    if not isinstance(subject, str) or not subject or not isinstance(issuer, str) or not issuer:
        raise ValueError("OIDC userinfo must contain non-empty 'sub' and 'iss' claims")

    display_name = next(
        (
            value
            for value in (
                userinfo.get("preferred_username"),
                userinfo.get("name"),
                userinfo.get("email"),
                subject,
            )
            if isinstance(value, str) and value
        ),
        subject,
    )
    now = datetime.now(timezone.utc)
    payload = {
        "sub": subject,
        "oidc_issuer": issuer,
        "name": display_name,
        "csrf": secrets.token_urlsafe(32),
        "iat": now,
        "exp": now + timedelta(minutes=settings.oidc_session_expire_minutes),
        "token_use": OIDC_SESSION_TOKEN_USE,
    }
    return str(jwt.encode(payload, settings.secret_key, algorithm=settings.algorithm))


def decode_oidc_session_token(token: str) -> OIDCSession | None:
    try:
        payload = jwt.decode(token, settings.secret_key, algorithms=[settings.algorithm])
        if payload.get("token_use") != OIDC_SESSION_TOKEN_USE:
            return None
        subject = payload.get("sub")
        issuer = payload.get("oidc_issuer")
        display_name = payload.get("name")
        csrf_token = payload.get("csrf")
        if not isinstance(subject, str) or not subject:
            return None
        if not isinstance(issuer, str) or not issuer:
            return None
        if not isinstance(display_name, str) or not display_name:
            return None
        if not isinstance(csrf_token, str) or not csrf_token:
            return None
        return OIDCSession(
            subject=subject,
            issuer=issuer,
            display_name=display_name,
            csrf_token=csrf_token,
        )
    except InvalidTokenError:
        return None


def oidc_session_cookie_name() -> str:
    return "__Host-homelable-session" if settings.oidc_cookie_secure else "homelable-session"


def set_oidc_session_cookie(response: Response, token: str) -> None:
    response.set_cookie(
        key=oidc_session_cookie_name(),
        value=token,
        max_age=settings.oidc_session_expire_minutes * 60,
        secure=settings.oidc_cookie_secure,
        httponly=True,
        samesite="lax",
        path="/",
    )


def clear_oidc_session_cookie(response: Response) -> None:
    response.delete_cookie(
        key=oidc_session_cookie_name(),
        secure=settings.oidc_cookie_secure,
        httponly=True,
        samesite="lax",
        path="/",
    )


def get_oidc_session(request: Request) -> OIDCSession | None:
    token = request.cookies.get(oidc_session_cookie_name())
    return decode_oidc_session_token(token) if token else None


def origin_is_allowed(origin: str | None) -> bool:
    if not origin:
        return False
    allowed_origins = {value.rstrip("/") for value in settings.cors_origins}
    # A same-origin deployment behind a reverse proxy needs no CORS, so
    # CORS_ORIGINS is often left at its localhost default — yet browsers still
    # send Origin on unsafe methods. The OIDC callback URL is served by this app,
    # so its origin is the app's own and is always a legitimate CSRF origin.
    app_origin = origin_of(settings.oidc_redirect_uri)
    if app_origin:
        allowed_origins.add(app_origin)
    return origin.rstrip("/") in allowed_origins


class OIDCCSRFMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        if settings.auth_mode != "oidc" or request.method in SAFE_METHODS:
            return await call_next(request)

        # Explicit non-cookie credentials take precedence in get_current_user,
        # so their requests keep the existing CSRF-free API behavior.
        if request.headers.get("Authorization") or request.headers.get("X-MCP-Service-Key"):
            return await call_next(request)

        oidc_session = get_oidc_session(request)
        if oidc_session is None:
            return await call_next(request)

        csrf_token = request.headers.get("X-Homelable-CSRF", "")
        csrf_valid = hmac.compare_digest(csrf_token.encode(), oidc_session.csrf_token.encode())
        if not csrf_valid or not origin_is_allowed(request.headers.get("Origin")):
            return JSONResponse(
                status_code=status.HTTP_403_FORBIDDEN,
                content={"detail": "CSRF validation failed"},
            )
        return await call_next(request)


@lru_cache(maxsize=8)
def _build_oidc_client(
    discovery_url: str,
    client_id: str,
    client_secret: str,
    scopes: str,
) -> StarletteOAuth2App:
    oauth = OAuth()
    client = oauth.register(
        name="oidc",
        client_id=client_id,
        client_secret=client_secret,
        server_metadata_url=discovery_url,
        client_kwargs={
            "scope": scopes,
            "code_challenge_method": "S256",
        },
    )
    if client is None:
        raise RuntimeError("Failed to register OIDC client")
    return client


def get_oidc_client() -> StarletteOAuth2App:
    return _build_oidc_client(
        settings.oidc_discovery_url,
        settings.oidc_client_id,
        settings.oidc_client_secret,
        settings.oidc_scopes,
    )