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 starlette.middleware.base import BaseHTTPMiddleware

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(BaseHTTPMiddleware):
    """Defense-in-depth: 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.
    """

    async def dispatch(self, request: Request, call_next):
        path = request.url.path
        if request.method == "OPTIONS":
            return await call_next(request)
        if _is_install_mode():
            return await call_next(request)
        if not path.startswith("/api/"):
            return await call_next(request)
        if _is_public_path(path):
            return await call_next(request)

        token = _extract_bearer_token(request)
        if not token:
            return _unauthorized_response("Missing Authorization header")

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

        request.state.admin = admin
        return await call_next(request)


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.
            token_tv = payload.get("tv", -1)
            if token_tv != admin.token_version:
                logger.debug(f"Token invalidated: token_version mismatch (token={token_tv}, db={admin.token_version})")
                return None

            # Secondary check: updated_at timestamp (defense in depth).
            # Catches any admin update not covered by token_version.
            token_updated = payload.get("updated", 0)
            if token_updated and admin.updated_at:
                try:
                    admin_updated_ts = int(admin.updated_at.timestamp())
                    if admin_updated_ts > token_updated + 5:  # 5s grace for clock skew
                        logger.debug(f"Token invalidated: admin updated after token issued")
                        return None
                except (ValueError, OSError):
                    pass
            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", -1)
        if token_tv != admin.token_version:
            return None

        # Secondary check: updated_at timestamp (defense in depth)
        token_updated = payload.get("updated", 0)
        if token_updated and admin.updated_at:
            try:
                admin_updated_ts = int(admin.updated_at.timestamp())
                if admin_updated_ts > token_updated + 5:  # 5s grace for clock skew
                    return None
            except (ValueError, OSError):
                pass

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