Nexus/server/api/auth_jwt.py

"""Nexus — JWT Authentication Middleware
Validates JWT access token on ALL protected API routes.

A1: JWT middleware deployed to all API routes.
Old API Key auth remains for /api/agent/* endpoints (Agent→Nexus communication).
"""

import logging
from pathlib import Path
from typing import Optional

_ROOT_DIR = Path(__file__).resolve().parent.parent.parent


def _is_install_mode() -> bool:
    """Check if the system is in install mode (no .env file).

    Evaluated at call time (not import time) so that the install wizard
    writing .env is detected immediately without requiring a restart.
    """
    return not (_ROOT_DIR / ".env").exists()

from fastapi import Depends, HTTPException, Request, status
from fastapi.responses import JSONResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

from server.domain.models import Admin
from server.api.dependencies import _get_request_session
from server.infrastructure.database.admin_repo import AdminRepositoryImpl
from server.config import settings

logger = logging.getLogger("nexus.auth_jwt")

security = HTTPBearer(auto_error=False)


# Routes that bypass JWT auth (Agent heartbeat uses API Key, login obviously doesn't need JWT)
PUBLIC_PREFIXES = (
    "/api/auth/login",
    "/api/auth/refresh",
    "/api/auth/logout",  # invalidates refresh token from body, no access JWT required
    "/api/agent/",       # Agent uses X-API-Key header
    "/api/install/",     # install wizard (403 when already installed, except /status)
    "/health",
    "/ws/",
    "/docs",
    "/openapi.json",
    "/redoc",
)


def _is_public_path(path: str) -> bool:
    """Check if the request path is public (doesn't require JWT)"""
    for prefix in PUBLIC_PREFIXES:
        if path.startswith(prefix):
            return True
    return False


def _unauthorized_response(detail: str) -> JSONResponse:
    return JSONResponse(
        status_code=status.HTTP_401_UNAUTHORIZED,
        content={"detail": detail},
        headers={"WWW-Authenticate": "Bearer"},
    )


def _extract_bearer_token(request: Request) -> Optional[str]:
    auth_header = request.headers.get("Authorization", "")
    if not auth_header.startswith("Bearer "):
        return None
    token = auth_header[7:].strip()
    return token or None


class JwtAuthMiddleware:
    """Pure ASGI middleware: reject unauthenticated /api/* before route handlers run.

    Must be registered inside DbSessionMiddleware so _verify_token can use request.state.db.
    Route-level Depends(get_current_admin) remains for injecting Admin into handlers.

    NOTE: Uses pure ASGI (not BaseHTTPMiddleware) to avoid MissingGreenlet errors
    with SQLAlchemy async. BaseHTTPMiddleware's call_next() runs the endpoint in
    a separate task, which breaks SQLAlchemy's async greenlet context.
    """

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        path = scope.get("path", "")

        # Skip OPTIONS (CORS preflight)
        # Read method from ASGI scope
        method = scope.get("method", "")
        if method == "OPTIONS":
            await self.app(scope, receive, send)
            return

        if _is_install_mode():
            await self.app(scope, receive, send)
            return

        if not path.startswith("/api/"):
            await self.app(scope, receive, send)
            return

        if _is_public_path(path):
            await self.app(scope, receive, send)
            return

        # Extract Bearer token from headers
        token = None
        for header_name, header_value in scope.get("headers", []):
            if header_name == b"authorization":
                auth_str = header_value.decode("latin-1", errors="replace")
                if auth_str.startswith("Bearer "):
                    token = auth_str[7:].strip()
                break

        if not token:
            response = _unauthorized_response("Missing Authorization header")
            await response(scope, receive, send)
            return

        # Verify token — need a Request object for _verify_token
        # Create a minimal Request from the ASGI scope
        from starlette.requests import Request
        request = Request(scope, receive, send)

        admin = await _verify_token(token, request)
        if not admin:
            response = _unauthorized_response("Invalid or expired token")
            await response(scope, receive, send)
            return

        # Store admin in scope state (accessible via request.state.admin)
        scope.setdefault("state", {})
        scope["state"]["admin"] = admin
        await self.app(scope, receive, send)


async def get_current_admin(
    request: Request,
    credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
) -> Optional[Admin]:
    """FastAPI dependency: extract JWT from Authorization header → return Admin or None

    A1: Applied to all API routes except public ones.
    For public routes, returns None if no credentials provided.
    For protected routes, raises 401 if no/invalid credentials.
    """
    # Skip JWT for public routes — return None (route handler decides what to do)
    if _is_public_path(request.url.path):
        if not credentials:
            return None

    cached = getattr(request.state, "admin", None)
    if cached is not None:
        return cached

    if not credentials:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing Authorization header",
            headers={"WWW-Authenticate": "Bearer"},
        )

    admin = await _verify_token(credentials.credentials, request)
    if not admin:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or expired token",
            headers={"WWW-Authenticate": "Bearer"},
        )

    return admin


async def _verify_token(token: str, request: Request) -> Optional[Admin]:
    """Verify JWT access token → return Admin or None"""
    try:
        import jwt as pyjwt
        payload = pyjwt.decode(
            token, settings.SECRET_KEY, algorithms=["HS256"],
            options={"require": ["exp", "sub"]},
        )
    except pyjwt.ExpiredSignatureError:
        logger.debug("JWT token expired")
        return None
    except pyjwt.InvalidTokenError as e:
        logger.debug(f"JWT invalid: {e}")
        return None
    except Exception:
        logger.debug("JWT verification unexpected error", exc_info=True)
        return None

    admin_id = payload.get("sub")
    if not admin_id:
        return None

    try:
        admin_id = int(admin_id)
    except (ValueError, TypeError):
        return None

    try:
        session = await _get_request_session(request)
        admin_repo = AdminRepositoryImpl(session)
        admin = await admin_repo.get_by_id(admin_id)
        if admin and admin.is_active:
            # Primary check: token_version must match exactly.
            # Changed on password change, TOTP disable, or token reuse detection.
            # This invalidates ALL access tokens immediately — no grace window.
            # Normalize None → 0 for safety (DB may have NULL from direct inserts)
            token_tv = payload.get("tv") or 0
            admin_tv = admin.token_version or 0
            if token_tv != admin_tv:
                logger.debug(f"Token invalidated: token_version mismatch (token={token_tv}, db={admin.token_version})")
                return None

            return admin
    except Exception:
        logger.debug("JWT admin lookup (get_current_admin) failed", exc_info=True)

    return None


async def get_optional_admin(
    request: Request,
    credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
) -> Optional[Admin]:
    """Optional JWT auth — returns Admin if token valid, None otherwise.
    Uses middleware session via request.state.db (no leaked sessions).
    """
    if not credentials:
        return None

    try:
        import jwt as pyjwt
        payload = pyjwt.decode(
            credentials.credentials, settings.SECRET_KEY, algorithms=["HS256"],
            options={"require": ["exp", "sub"]},
        )
        admin_id = payload.get("sub")
        if not admin_id:
            return None
        try:
            admin_id = int(admin_id)
        except (ValueError, TypeError):
            return None

        # Use middleware session instead of creating independent session
        session = await _get_request_session(request)
        admin_repo = AdminRepositoryImpl(session)
        admin = await admin_repo.get_by_id(admin_id)
        if not admin or not admin.is_active:
            return None

        # Primary check: token_version must match exactly
        token_tv = payload.get("tv") or 0
        admin_tv = admin.token_version or 0
        if token_tv != admin_tv:
            return None

        return admin
    except Exception:
        logger.debug("JWT admin lookup failed", exc_info=True)
        return None