vignette/apps/api/app/routes/auth.py

"""BFF authentication routes.

The production path is Google OIDC authorization code + PKCE. Until the DB
session table is wired, the issued browser sessions are server-side in-proc
sessions backed by an opaque HttpOnly cookie. Local development also has a
dev-only server login endpoint so Playwright can exercise auth without trusting
browser localStorage.
"""

from __future__ import annotations

import base64
import hashlib
import secrets
import time
from dataclasses import dataclass
from typing import Annotated, Literal, Optional
from urllib.parse import urlencode, urlsplit

import httpx
from fastapi import APIRouter, Cookie, HTTPException, Query, Request, Response, status
from fastapi.responses import RedirectResponse
from pydantic import BaseModel

from ..auth_sessions import InactiveUserError, SessionUser, create_session, revoke_session
from ..config import settings
from ..deps import CurrentPrincipal, Principal, Role
from ..saml import (
    SamlIdentity,
    acs_url_for_entity_id,
    build_authn_request,
    parse_fixture_response,
    redirect_binding_url,
)

router = APIRouter(prefix="/auth", tags=["auth"])

GOOGLE_AUTHORIZE_URL = "https://accounts.google.com/o/oauth2/v2/auth"
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
GOOGLE_TOKENINFO_URL = "https://oauth2.googleapis.com/tokeninfo"
OAUTH_STATE_TTL_SECONDS = 10 * 60


@dataclass(slots=True)
class OAuthState:
    code_verifier: str
    next_path: str
    created_at: float


@dataclass(slots=True)
class SamlState:
    request_id: str
    next_path: str
    created_at: float


_oauth_states: dict[str, OAuthState] = {}
_saml_states: dict[str, SamlState] = {}


class MeResponse(BaseModel):
    user_id: str
    email: str
    display_name: str
    role: str
    cohort_ids: list[str]


class AuthProviderStatus(BaseModel):
    provider: Literal["google", "saml"]
    configured: bool
    enabled: bool
    login_path: str


class AuthConfigResponse(BaseModel):
    google_oauth_configured: bool
    saml_configured: bool
    providers: list[AuthProviderStatus]
    allowed_email_domains: list[str]
    redirect_uri: str
    dev_login_enabled: bool


class DevLoginRequest(BaseModel):
    email: str
    role: Literal["learner", "teacher", "admin"] = "learner"
    display_name: str | None = None


def _normalize_domain(domain: str | None) -> str:
    return (domain or "").strip().lower().lstrip("@")


def _normalize_email(email: str | None) -> str:
    return (email or "").strip().lower()


def _email_domain(email: str | None) -> str:
    value = _normalize_email(email)
    if "@" not in value:
        return ""
    return value.rsplit("@", 1)[1]


def _normalize_email_set(values: list[str]) -> set[str]:
    return {email for value in values if (email := _normalize_email(value))}


def _google_configured() -> bool:
    return bool(settings.oauth_google_client_id and settings.oauth_google_client_secret)


def _saml_configured() -> bool:
    return bool(
        settings.auth_saml_enabled
        and settings.saml_sp_entity_id.strip()
        and settings.saml_sso_url.strip()
    )


def _auth_provider_statuses() -> list[AuthProviderStatus]:
    google_ready = _google_configured()
    saml_ready = _saml_configured()
    return [
        AuthProviderStatus(
            provider="google",
            configured=google_ready,
            enabled=google_ready,
            login_path="/auth/login?provider=google",
        ),
        AuthProviderStatus(
            provider="saml",
            configured=saml_ready,
            enabled=saml_ready,
            login_path="/auth/login?provider=saml",
        ),
    ]


def allowed_email_domains() -> set[str]:
    """Configured login email domains, normalized for claim checks."""
    return {
        normalized
        for domain in settings.auth_allowed_email_domains
        if (normalized := _normalize_domain(domain))
    }


def validate_google_identity_domain(
    *,
    email: str | None,
    email_verified: bool,
    hosted_domain: str | None = None,
) -> str:
    """Reject Google identities outside the allowed email domain list.

    Google Console authorized domains protect app/redirect domains, not user
    email domains. After id_token signature/audience/issuer validation, call
    this check with the `email`, `email_verified`, and optional `hd` claims.
    """
    normalized_email = _normalize_email(email)
    domain = _email_domain(normalized_email)
    allowed = allowed_email_domains()

    if not allowed:
        raise HTTPException(
            status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="allowed email domains are not configured",
        )
    if not normalized_email or not domain:
        raise HTTPException(status.HTTP_403_FORBIDDEN, detail="email claim is required")
    if not email_verified:
        raise HTTPException(status.HTTP_403_FORBIDDEN, detail="email is not verified")
    if domain not in allowed:
        raise HTTPException(status.HTTP_403_FORBIDDEN, detail="email domain is not allowed")

    hd = _normalize_domain(hosted_domain)
    if hd and hd not in allowed:
        raise HTTPException(status.HTTP_403_FORBIDDEN, detail="hosted domain is not allowed")

    return normalized_email


def _role_for_email(email: str) -> Role:
    normalized = _normalize_email(email)
    if normalized in _normalize_email_set(settings.auth_admin_emails):
        return Role.ADMIN
    if normalized in _normalize_email_set(settings.auth_teacher_emails):
        return Role.TEACHER
    return Role.LEARNER


def _role_for_saml_identity(identity: SamlIdentity) -> Role:
    hinted = (identity.role_hint or "").strip().lower()
    if hinted in {"admin", "administrator"}:
        return Role.ADMIN
    if hinted in {"teacher", "instructor", "faculty"}:
        return Role.TEACHER
    return _role_for_email(identity.email)


def _safe_next_path(next_path: str | None) -> str:
    if not next_path or not next_path.startswith("/") or next_path.startswith("//"):
        return "/"
    return next_path


def _url_origin(value: str | None) -> str | None:
    if not value:
        return None
    parsed = urlsplit(value)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        return None
    return f"{parsed.scheme}://{parsed.netloc}".rstrip("/")


def _is_local_origin(origin: str) -> bool:
    host = urlsplit(origin).hostname or ""
    return host in {"localhost", "127.0.0.1", "::1"}


def _configured_frontend_origins() -> list[str]:
    origins: list[str] = []
    for value in [settings.frontend_base_url, *settings.cors_origins]:
        origin = _url_origin(value)
        if origin and origin not in origins:
            origins.append(origin)
    return origins


def _frontend_origin_for_request(request: Request | None = None) -> str:
    origins = _configured_frontend_origins()
    fallback = (_url_origin(settings.frontend_base_url) or "http://localhost:5173").rstrip("/")

    if request is not None:
        for header_name in ("origin", "referer"):
            candidate = _url_origin(request.headers.get(header_name))
            if candidate in origins:
                return candidate

        forwarded_host = request.headers.get("x-forwarded-host")
        host = (forwarded_host or request.headers.get("host") or "").split(",", 1)[0].strip()
        hostname = host.rsplit(":", 1)[0].lower() if host else ""
        if hostname == "api-vignette.chanpaca.net":
            return "https://vignette.chanpaca.net"
        if hostname in {"localhost", "127.0.0.1", "::1"}:
            return fallback

    if not _is_local_origin(fallback):
        return fallback

    for origin in origins:
        hostname = (urlsplit(origin).hostname or "").lower()
        if not _is_local_origin(origin) and hostname != "api-vignette.chanpaca.net":
            return origin
    return fallback


def _frontend_url(path: str, request: Request | None = None) -> str:
    base = _frontend_origin_for_request(request)
    return f"{base}{_safe_next_path(path)}"


def _pkce_challenge(verifier: str) -> str:
    digest = hashlib.sha256(verifier.encode("ascii")).digest()
    return base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")


def _prune_oauth_states() -> None:
    cutoff = time.time() - OAUTH_STATE_TTL_SECONDS
    stale = [key for key, value in _oauth_states.items() if value.created_at < cutoff]
    for key in stale:
        _oauth_states.pop(key, None)


def _prune_saml_states() -> None:
    cutoff = time.time() - OAUTH_STATE_TTL_SECONDS
    stale = [key for key, value in _saml_states.items() if value.created_at < cutoff]
    for key in stale:
        _saml_states.pop(key, None)


def _cookie_secure() -> bool:
    # The __Host- prefix requires Secure, Path=/, and no Domain. Modern Chrome
    # accepts Secure cookies on localhost, which keeps dev and prod semantics
    # aligned.
    return settings.is_prod or settings.cookie_name.startswith("__Host-")


def _set_session_cookie(response: Response, sid: str) -> None:
    response.set_cookie(
        key=settings.cookie_name,
        value=sid,
        max_age=settings.session_ttl_seconds,
        httponly=True,
        secure=_cookie_secure(),
        samesite="lax",
        path="/",
    )
    if settings.environment == "dev":
        response.set_cookie(
            key="vignette_sid",
            value=sid,
            max_age=settings.session_ttl_seconds,
            httponly=True,
            secure=False,
            samesite="lax",
            path="/",
        )


def _delete_session_cookie(response: Response) -> None:
    response.delete_cookie(
        settings.cookie_name,
        httponly=True,
        secure=_cookie_secure(),
        samesite="lax",
        path="/",
    )
    if settings.environment == "dev":
        response.delete_cookie(
            "vignette_sid",
            httponly=True,
            secure=False,
            samesite="lax",
            path="/",
        )


def _me_response(user: SessionUser | Principal) -> MeResponse:
    return MeResponse(
        user_id=user.user_id,
        email=getattr(user, "email", ""),
        display_name=getattr(user, "display_name", "") or getattr(user, "email", ""),
        role=user.role.value if isinstance(user.role, Role) else user.role,
        cohort_ids=user.cohort_ids,
    )


def _frontend_login_redirect(reason: str, request: Request) -> RedirectResponse:
    base_url = _frontend_origin_for_request(request)
    return RedirectResponse(f"{base_url}/login?{urlencode({'oauth': reason})}", status_code=302)


def _dev_login_available(request: Request) -> bool:
    if settings.environment != "dev" or not settings.auth_dev_login_enabled:
        return False

    for header_name in ("origin", "referer"):
        origin = _url_origin(request.headers.get(header_name))
        if origin and not _is_local_origin(origin):
            return False

    forwarded_host = request.headers.get("x-forwarded-host")
    host = (forwarded_host or request.headers.get("host") or "").split(",", 1)[0].strip()
    origin = _url_origin(f"http://{host}") if host else None
    return bool(origin and _is_local_origin(origin))


@router.get("/config", response_model=AuthConfigResponse)
async def auth_config(request: Request) -> AuthConfigResponse:
    """Return non-secret login configuration for the browser login screen."""
    google_ready = _google_configured()
    saml_ready = _saml_configured()
    return AuthConfigResponse(
        google_oauth_configured=google_ready,
        saml_configured=saml_ready,
        providers=_auth_provider_statuses(),
        allowed_email_domains=sorted(allowed_email_domains()),
        redirect_uri=settings.oauth_redirect_uri,
        dev_login_enabled=_dev_login_available(request),
    )


@router.get("/login")
async def login(
    request: Request,
    provider: Annotated[str, Query()] = "google",
    next: Annotated[str | None, Query()] = None,
) -> RedirectResponse:
    """Start Google OIDC authorization code + PKCE login."""
    if provider == "saml":
        if not _saml_configured():
            return _frontend_login_redirect("saml_not_configured", request)
        _prune_saml_states()
        relay_state = secrets.token_urlsafe(32)
        acs_url = acs_url_for_entity_id(settings.saml_sp_entity_id)
        request_id, authn_request_xml = build_authn_request(
            sp_entity_id=settings.saml_sp_entity_id,
            sso_url=settings.saml_sso_url,
            acs_url=acs_url,
        )
        _saml_states[relay_state] = SamlState(
            request_id=request_id,
            next_path=_safe_next_path(next),
            created_at=time.time(),
        )
        return RedirectResponse(
            redirect_binding_url(
                sso_url=settings.saml_sso_url,
                authn_request_xml=authn_request_xml,
                relay_state=relay_state,
            ),
            status_code=302,
        )
    if provider != "google":
        return _frontend_login_redirect("unsupported_provider", request)
    if not _google_configured():
        return _frontend_login_redirect("not_configured", request)

    _prune_oauth_states()
    state = secrets.token_urlsafe(32)
    verifier = secrets.token_urlsafe(64)
    _oauth_states[state] = OAuthState(
        code_verifier=verifier,
        next_path=_safe_next_path(next),
        created_at=time.time(),
    )

    params = {
        "client_id": settings.oauth_google_client_id,
        "redirect_uri": settings.oauth_redirect_uri,
        "response_type": "code",
        "scope": "openid email profile",
        "state": state,
        "code_challenge": _pkce_challenge(verifier),
        "code_challenge_method": "S256",
        "prompt": "select_account",
    }
    return RedirectResponse(f"{GOOGLE_AUTHORIZE_URL}?{urlencode(params)}", status_code=302)


@router.get("/callback")
async def callback(
    request: Request,
    code: Annotated[Optional[str], Query()] = None,
    state: Annotated[Optional[str], Query()] = None,
) -> RedirectResponse:
    """Exchange Google auth code, validate identity, and issue a BFF cookie."""
    if not code or not state:
        return _frontend_login_redirect("missing_callback", request)

    _prune_oauth_states()
    stored = _oauth_states.pop(state, None)
    if stored is None:
        return _frontend_login_redirect("invalid_state", request)

    async with httpx.AsyncClient(timeout=10.0) as client:
        token_res = await client.post(
            GOOGLE_TOKEN_URL,
            data={
                "client_id": settings.oauth_google_client_id,
                "client_secret": settings.oauth_google_client_secret,
                "code": code,
                "grant_type": "authorization_code",
                "redirect_uri": settings.oauth_redirect_uri,
                "code_verifier": stored.code_verifier,
            },
            headers={"Accept": "application/json"},
        )
        if token_res.status_code >= 400:
            return _frontend_login_redirect("token_exchange_failed", request)
        token_payload = token_res.json()
        id_token = token_payload.get("id_token")
        if not isinstance(id_token, str) or not id_token:
            return _frontend_login_redirect("id_token_missing", request)

        info_res = await client.get(GOOGLE_TOKENINFO_URL, params={"id_token": id_token})
        if info_res.status_code >= 400:
            return _frontend_login_redirect("id_token_invalid", request)
        claims = info_res.json()

    if claims.get("aud") != settings.oauth_google_client_id:
        return _frontend_login_redirect("audience_mismatch", request)
    issuer = claims.get("iss")
    if issuer not in {"accounts.google.com", "https://accounts.google.com"}:
        return _frontend_login_redirect("issuer_mismatch", request)

    try:
        email = validate_google_identity_domain(
            email=claims.get("email"),
            email_verified=claims.get("email_verified") in {True, "true", "True", "1", 1},
            hosted_domain=claims.get("hd"),
        )
    except HTTPException:
        return _frontend_login_redirect("domain_not_allowed", request)
    role = _role_for_email(email)
    display_name = str(claims.get("name") or email)
    try:
        sid, _ = await create_session(
            email=email,
            display_name=display_name,
            role=role.value,
            cohort_ids=[],
        )
    except InactiveUserError as exc:
        return _frontend_login_redirect("inactive_user", request)

    response = RedirectResponse(_frontend_url(stored.next_path, request), status_code=302)
    _set_session_cookie(response, sid)
    return response


@router.post("/saml/acs")
async def saml_acs(request: Request) -> RedirectResponse:
    """Accept a minimal unsigned SAMLResponse for local fixture SAML proof.

    Signed SAML verification is intentionally not implemented. When
    SAML_X509_CERT_FINGERPRINT is configured, this endpoint refuses to trust the
    response so production does not silently run unsigned SAML.
    """
    if not _saml_configured():
        return _frontend_login_redirect("saml_not_configured", request)
    if settings.saml_x509_cert_fingerprint.strip():
        return _frontend_login_redirect("saml_signature_verification_required", request)
    if settings.environment != "dev":
        return _frontend_login_redirect("saml_fixture_acs_dev_only", request)

    form = await request.form()
    relay_state = str(form.get("RelayState") or "")
    encoded_response = str(form.get("SAMLResponse") or "")
    if not relay_state or not encoded_response:
        return _frontend_login_redirect("saml_missing_callback", request)

    _prune_saml_states()
    stored = _saml_states.pop(relay_state, None)
    if stored is None:
        return _frontend_login_redirect("saml_invalid_state", request)

    try:
        identity = parse_fixture_response(encoded_response)
        email = validate_google_identity_domain(
            email=identity.email,
            email_verified=True,
            hosted_domain=_email_domain(identity.email),
        )
    except (HTTPException, ValueError):
        return _frontend_login_redirect("saml_assertion_invalid", request)

    role = _role_for_saml_identity(identity)
    try:
        sid, _ = await create_session(
            email=email,
            display_name=identity.display_name or email,
            role=role.value,
            cohort_ids=[],
        )
    except InactiveUserError:
        return _frontend_login_redirect("inactive_user", request)

    response = RedirectResponse(_frontend_url(stored.next_path, request), status_code=302)
    _set_session_cookie(response, sid)
    return response


@router.post("/dev-login", response_model=MeResponse)
async def dev_login(request: Request, body: DevLoginRequest, response: Response) -> MeResponse:
    """Dev-only server login for local E2E and manual testing.

    This is not a browser-side auth shortcut: the role is stored server-side and
    the browser only gets the same opaque HttpOnly cookie used by OAuth.
    """
    if not _dev_login_available(request):
        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="dev login is disabled")

    email = validate_google_identity_domain(
        email=str(body.email),
        email_verified=True,
        hosted_domain=_email_domain(str(body.email)),
    )
    try:
        sid, user = await create_session(
            email=email,
            display_name=body.display_name or email,
            role=body.role,
            cohort_ids=[],
        )
    except InactiveUserError as exc:
        raise HTTPException(status.HTTP_403_FORBIDDEN, detail="user is inactive") from exc
    _set_session_cookie(response, sid)
    return _me_response(user)


@router.post("/logout")
async def logout(
    response: Response,
    session_cookie: Annotated[Optional[str], Cookie(alias="__Host-vignette_sid")] = None,
    dev_session_cookie: Annotated[Optional[str], Cookie(alias="vignette_sid")] = None,
) -> dict[str, bool]:
    """Revoke the current server session and expire the browser cookie."""
    await revoke_session(session_cookie or (dev_session_cookie if settings.environment == "dev" else None))
    _delete_session_cookie(response)
    return {"ok": True}


@router.get("/me", response_model=MeResponse)
async def me(principal: CurrentPrincipal) -> MeResponse:
    """Return the current authenticated user. Unauthenticated requests are 401."""
    return _me_response(principal)