Aegis/backend/app/routers/sso.py

"""Phase 14: SSO / SAML 2.0 router."""

import os
from urllib.parse import urlparse

from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import RedirectResponse
from sqlalchemy.orm import Session

from app.database import get_db
from app.dependencies.auth import require_any_role
from app import auth as auth_lib
from app.schemas.sso_schema import (
    SsoConfigCreate, SsoConfigOut, SsoStatusResponse,
)
import app.services.sso_service as svc

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

_COOKIE_NAME = "aegis_token"

# Mirror the same SECURE_COOKIES logic used in the auth router so that
# SAML-authenticated sessions respect the deployment's HTTPS configuration.
_aegis_env = os.environ.get("AEGIS_ENV", "development").lower()
_secure_cookie_env = os.environ.get("SECURE_COOKIES", "auto").lower()
if _secure_cookie_env == "false":
    _IS_HTTPS = False
elif _secure_cookie_env == "true":
    _IS_HTTPS = True
else:  # "auto" — active only when AEGIS_ENV=production
    _IS_HTTPS = _aegis_env == "production"

_COOKIE_OPTS = {"httponly": True, "samesite": "lax", "secure": _IS_HTTPS}


# ── Public ────────────────────────────────────────────────────────────────────

@router.get("/status", response_model=SsoStatusResponse)
def sso_status(db: Session = Depends(get_db)):
    """Return whether SSO is enabled and configured (public — for login page)."""
    return svc.get_status(db)


@router.get("/metadata", response_model=None)
def sp_metadata(db: Session = Depends(get_db)):
    """
    Return the Service Provider SAML metadata XML.

    Upload this XML to your IdP (Okta, Azure AD, etc.) to register Aegis.
    """
    try:
        xml = svc.get_sp_metadata(db)
    except Exception as exc:
        raise HTTPException(status_code=503, detail=str(exc))
    return Response(content=xml, media_type="application/xml")


@router.get("/login")
def sso_login(request: Request, db: Session = Depends(get_db)):
    """
    Initiate SAML login — redirects the browser to the IdP.

    The IdP will POST the SAML Response to ``/sso/callback`` after authentication.
    """
    request_data = {
        "https":        request.url.scheme == "https",
        "http_host":    request.url.hostname,
        "path":         request.url.path,
        "port":         str(request.url.port or (443 if request.url.scheme == "https" else 80)),
        "get_data":     dict(request.query_params),
        "post_data":    {},
        "query_string": str(request.url.query),
    }
    try:
        result = svc.initiate_login(db, request_data)
    except RuntimeError as exc:
        raise HTTPException(status_code=503, detail=str(exc))
    redirect_url = result["redirect_url"]
    if urlparse(redirect_url).scheme not in ("http", "https"):
        raise HTTPException(status_code=400, detail="Invalid IdP redirect URL")
    return RedirectResponse(url=redirect_url)


@router.post("/callback")
async def sso_callback(request: Request, db: Session = Depends(get_db)):
    """
    SAML Assertion Consumer Service (ACS) endpoint.

    The IdP POSTs the SAML Response here.  On success, sets the aegis_token
    cookie and redirects to the frontend.
    """
    form = await request.form()
    request_data = {
        "https":        request.url.scheme == "https",
        "http_host":    request.url.hostname,
        "path":         request.url.path,
        "port":         str(request.url.port or (443 if request.url.scheme == "https" else 80)),
        "get_data":     dict(request.query_params),
        "post_data":    dict(form),
        "query_string": str(request.url.query),
    }
    try:
        user = svc.process_callback(db, request_data)
    except (ValueError, RuntimeError) as exc:
        raise HTTPException(status_code=401, detail=str(exc))

    access_token = auth_lib.create_access_token({"sub": user.username})
    response = RedirectResponse(url="/", status_code=302)
    response.set_cookie(_COOKIE_NAME, access_token, **_COOKIE_OPTS)
    return response


# ── Admin configuration ────────────────────────────────────────────────────────

@router.get("/config", response_model=SsoConfigOut)
def get_sso_config(
    db: Session = Depends(get_db),
    _user=Depends(require_any_role("admin")),
):
    """Return the current SSO configuration (admin only)."""
    cfg = svc.get_config(db)
    if not cfg:
        raise HTTPException(status_code=404, detail="SSO not configured yet")
    return SsoConfigOut.model_validate(cfg)


@router.put("/config", response_model=SsoConfigOut)
def upsert_sso_config(
    body: SsoConfigCreate,
    db: Session = Depends(get_db),
    _user=Depends(require_any_role("admin")),
):
    """Create or replace the SSO configuration (admin only)."""
    cfg = svc.upsert_config(db, **body.model_dump(exclude_unset=False))
    return SsoConfigOut.model_validate(cfg)