Aegis/backend/app/dependencies/auth.py

"""
Authentication and RBAC dependencies for FastAPI.

Provides:
- ``get_current_user``: decodes JWT from HttpOnly cookie (preferred) or
  Authorization header (fallback), fetches user from DB, raises 401 on failure.
  Also accepts Aegis API keys (``aegis_…`` prefix) as Bearer tokens.
- ``require_role``: factory that returns a dependency enforcing a specific role
  (admins always pass).
"""

from typing import Optional

from fastapi import Cookie, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt
from sqlalchemy.orm import Session

from app import auth as auth_lib
from app.config import settings
from app.database import get_db
from app.models.user import User
from app.models.api_key import KEY_PREFIX

# ---------------------------------------------------------------------------
# OAuth2 scheme (reads Authorization header — used as fallback / Swagger UI)
# ---------------------------------------------------------------------------

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/v1/auth/login", auto_error=False)

# Cookie name — must match the one set in the auth router
_COOKIE_NAME = "aegis_token"

# ---------------------------------------------------------------------------
# Current-user dependency
# ---------------------------------------------------------------------------


async def get_current_user(
    aegis_token: Optional[str] = Cookie(None),
    bearer_token: Optional[str] = Depends(oauth2_scheme),
    db: Session = Depends(get_db),
) -> User:
    """Decode the JWT, look up the user in *db*, and return it.

    Token resolution order:
    1. ``aegis_token`` **HttpOnly cookie** (preferred — immune to XSS).
    2. ``Authorization: Bearer <token>`` header (fallback for API clients
       and Swagger UI).

    Raises :class:`~fastapi.HTTPException` **401** when:
    - no token is found in either location,
    - the token cannot be decoded,
    - the ``sub`` claim is missing, or
    - no matching active user exists in the database.
    """
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    revoked_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Token has been revoked",
        headers={"WWW-Authenticate": "Bearer"},
    )

    # Prefer cookie, fall back to header
    token = aegis_token or bearer_token
    if token is None:
        raise credentials_exception

    # ── API Key path (Bearer token starts with "aegis_") ──────────────────
    if token.startswith(KEY_PREFIX):
        from app.services.api_key_service import authenticate_raw_key
        user = authenticate_raw_key(db, token)
        if user is None:
            raise credentials_exception
        return user

    # ── JWT path ──────────────────────────────────────────────────────────
    try:
        payload = jwt.decode(
            token,
            settings.SECRET_KEY,
            algorithms=[settings.ALGORITHM],
        )
        username: str | None = payload.get("sub")
        if username is None:
            raise credentials_exception
        # Check token blacklist (revoked tokens)
        jti: str | None = payload.get("jti")
        if jti and auth_lib.is_token_blacklisted(jti):
            raise revoked_exception
    except JWTError:
        raise credentials_exception

    user = db.query(User).filter(User.username == username).first()
    if user is None or not user.is_active:
        raise credentials_exception

    return user


# ---------------------------------------------------------------------------
# Role-based access control dependency
# ---------------------------------------------------------------------------


async def require_password_changed(
    current_user: User = Depends(get_current_user),
) -> User:
    """Block all requests when the user still needs to change their password.

    Only ``/auth/change-password`` and ``/auth/me`` are exempt — those
    endpoints do **not** depend on this function.
    """
    if getattr(current_user, "must_change_password", False):
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="PASSWORD_CHANGE_REQUIRED",
        )
    return current_user


def require_role(required_role: str):
    """Return a FastAPI dependency that enforces *required_role*.

    The dependency allows the request to proceed when
    ``user.role == required_role`` **or** ``user.role == "admin"``.
    Otherwise it raises :class:`~fastapi.HTTPException` **403**.
    """

    async def role_checker(
        current_user: User = Depends(get_current_user),
    ) -> User:
        if current_user.role != required_role and current_user.role != "admin":
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Not enough permissions",
            )
        return current_user

    return role_checker


def require_any_role(*roles: str):
    """Return a FastAPI dependency that enforces **any** of the given *roles*.

    Admins always pass.  Usage example::

        @router.patch("/resource", dependencies=[Depends(require_any_role("red_lead", "blue_lead"))])
    """

    async def role_checker(
        current_user: User = Depends(get_current_user),
    ) -> User:
        if current_user.role != "admin" and current_user.role not in roles:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Not enough permissions",
            )
        return current_user

    return role_checker