SupplyAgent/api/auth_middleware.py

from __future__ import annotations

import re
from collections.abc import Awaitable, Callable

from fastapi import Request
from starlette.concurrency import run_in_threadpool
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse, Response

from api.services.auth import ADMIN_ROLE, SESSION_COOKIE_NAME, resolve_session

_PUBLIC_PATHS = {
    "/api/auth/login",
}
_AUTHENTICATED_USER_PATHS = {
    ("GET", "/api/auth/me"),
    ("POST", "/api/auth/logout"),
}
_NORMAL_USER_PATHS = {
    ("GET", "/api/category-tree"),
    ("GET", "/api/growth-category-tree"),
    ("GET", "/api/demand-grade"),
    ("GET", "/api/video-discovery/demands"),
    ("GET", "/api/video-discovery/runs"),
    ("GET", "/api/video-discovery/feedback"),
    ("POST", "/api/video-discovery/feedback"),
}
_NORMAL_USER_PATTERNS = (
    re.compile(r"^/api/video-discovery/demands/\d+$"),
    re.compile(r"^/api/video-discovery/runs/[^/]+$"),
    re.compile(r"^/api/video-discovery/runs/[^/]+/searches$"),
    re.compile(r"^/api/video-discovery/runs/[^/]+/candidates$"),
    re.compile(r"^/api/demand-grade/\d+/videos$"),
)


def normal_user_can_access(method: str, path: str) -> bool:
    if (method, path) in _AUTHENTICATED_USER_PATHS:
        return True
    if (method, path) in _NORMAL_USER_PATHS:
        return True
    return method == "GET" and any(pattern.fullmatch(path) for pattern in _NORMAL_USER_PATTERNS)


def _is_protected_path(path: str) -> bool:
    return (
        path == "/api"
        or path.startswith("/api/")
        or path in {"/docs", "/redoc", "/openapi.json"}
    )


class AuthenticationMiddleware(BaseHTTPMiddleware):
    """Authenticate API requests and enforce the two fixed application roles."""

    async def dispatch(
        self,
        request: Request,
        call_next: Callable[[Request], Awaitable[Response]],
    ) -> Response:
        path = request.url.path
        method = request.method.upper()
        if method == "OPTIONS" or not _is_protected_path(path) or path in _PUBLIC_PATHS:
            return await call_next(request)

        token = request.cookies.get(SESSION_COOKIE_NAME)
        user = await run_in_threadpool(resolve_session, token) if token else None
        if user is None:
            return JSONResponse(
                status_code=401,
                content={"detail": "authentication required"},
            )

        request.state.current_user = user
        if user["role"] == ADMIN_ROLE or normal_user_can_access(method, path):
            return await call_next(request)
        return JSONResponse(
            status_code=403,
            content={"detail": "permission denied"},
        )