app/api/sso.py

"""Author RAG — SSO / SAML 2.0 Integration.

Allows Enterprise authors (agency plan) to log in with their
company's identity provider (Okta, Azure AD, Google Workspace, etc.)

Flow:
  1. Agency admin configures their IdP metadata URL in admin panel
  2. Author visits /sso/{agency_slug} → redirected to IdP login
  3. IdP posts SAML assertion to /sso/{agency_slug}/acs
  4. We verify the assertion, find/create the User, issue a JWT
  5. Browser redirected to /admin with JWT in query param

Dependencies:
  pip install python3-saml  (wraps OpenSSL + lxml)

Routes:
  GET  /sso/{agency_slug}        → SAML AuthnRequest (redirect to IdP)
  POST /sso/{agency_slug}/acs    → SAML Assertion Consumer Service
  GET  /sso/{agency_slug}/metadata → SP metadata XML (upload to IdP)
  PUT  /api/admin/{slug}/sso/config → save IdP metadata URL (admin only)
"""

import structlog
from fastapi import APIRouter, Depends, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, RedirectResponse, Response
from pydantic import BaseModel, Field, HttpUrl
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import get_settings
from app.dependencies import get_admin_author_access, get_db

logger = structlog.get_logger(__name__)
cfg = get_settings()
router = APIRouter()

# In-memory SSO config store (Phase 4: encrypted DB column)
_sso_configs: dict[str, dict] = {}


class SSOConfig(BaseModel):
    idp_metadata_url: HttpUrl = Field(..., description="IdP metadata XML URL (Okta/Azure/Google)")
    attribute_email: str = Field("email", description="SAML attribute name for user email")
    attribute_name: str = Field("displayName", description="SAML attribute name for full name")
    auto_provision: bool = Field(True, description="Auto-create user accounts on first SSO login")


# ── Admin config ─────────────────────────────────────────────

@router.put("/{slug}/sso/config")
async def configure_sso(
    slug: str, body: SSOConfig,
    db: AsyncSession = Depends(get_db),
    access=Depends(get_admin_author_access),
):
    """Save IdP metadata URL for SSO."""
    author = access.author

    from app.models.pricing_plan import PricingPlan
    if not access.is_superadmin_proxy:
        plan = await db.get(PricingPlan, getattr(author, "stripe_plan_id", "") or "")
        if not plan or not getattr(plan, "has_white_label", False):
            raise HTTPException(403, "SSO requires Enterprise plan.")

    _sso_configs[slug] = body.model_dump()
    logger.info("SSO configured", author_id=slug, idp=str(body.idp_metadata_url))
    return {"message": "SSO configured. Test by visiting /sso/" + slug}


@router.get("/{slug}/sso/config")
async def get_sso_config(slug: str, access=Depends(get_admin_author_access)):
    cfg_data = _sso_configs.get(slug)
    if not cfg_data: return {"configured": False}
    return {"configured": True, **{k: v for k, v in cfg_data.items() if k != "idp_metadata_url"},
            "idp_metadata_url": "***configured***"}


# ── SAML SP endpoints ─────────────────────────────────────────

@router.get("/sso/{agency_slug}", tags=["SSO"])
async def sso_login(agency_slug: str, request: Request):
    """Redirect browser to IdP login page (SAML AuthnRequest)."""
    sso_cfg = _sso_configs.get(agency_slug)
    if not sso_cfg:
        raise HTTPException(404, f"No SSO configured for '{agency_slug}'.")

    try:
        from onelogin.saml2.auth import OneLogin_Saml2_Auth
        saml = _build_saml_auth(request, agency_slug, sso_cfg)
        redirect_url = saml.login()
        return RedirectResponse(url=redirect_url)
    except ImportError:
        # python3-saml not installed — return informative error
        return HTMLResponse(
            "<h2>SSO not available</h2>"
            "<p>Install: <code>pip install python3-saml</code></p>",
            status_code=503,
        )


@router.post("/sso/{agency_slug}/acs", tags=["SSO"])
async def sso_acs(
    agency_slug: str, request: Request,
    SAMLResponse: str = Form(...),
    db: AsyncSession = Depends(get_db),
):
    """SAML Assertion Consumer Service — validates assertion, issues JWT."""
    sso_cfg = _sso_configs.get(agency_slug)
    if not sso_cfg:
        raise HTTPException(404, "SSO not configured.")

    try:
        from onelogin.saml2.auth import OneLogin_Saml2_Auth
        saml = _build_saml_auth(request, agency_slug, sso_cfg)
        saml.process_response()
        errors = saml.get_errors()
        if errors:
            raise HTTPException(400, f"SAML errors: {', '.join(errors)}")

        attrs = saml.get_attributes()
        email = _first(attrs.get(sso_cfg.get("attribute_email", "email"), []))
        name = _first(attrs.get(sso_cfg.get("attribute_name", "displayName"), []))

        if not email:
            raise HTTPException(400, "IdP did not return an email attribute.")

        # Find or auto-provision user
        from sqlalchemy import select
        from app.models.user import User
        user_r = await db.execute(select(User).where(User.email == email))
        user = user_r.scalar_one_or_none()

        if not user:
            if not sso_cfg.get("auto_provision", True):
                raise HTTPException(403, "User not found and auto-provisioning is disabled.")
            import bcrypt, secrets
            user = User(
                email=email, full_name=name or email.split("@")[0],
                password_hash=bcrypt.hashpw(secrets.token_bytes(32), bcrypt.gensalt()).decode(),
                role="author", is_active=True,
                stripe_plan_id=None,  # Agency plan inherited from owner
            )
            db.add(user)
            await db.commit()
            logger.info("SSO auto-provisioned user", email=email, agency=agency_slug)

        from app.services.auth_service import create_access_token
        token = create_access_token({"sub": user.id, "email": user.email, "sso": agency_slug})

        # Redirect to admin with token in URL fragment (never hits server logs)
        return RedirectResponse(
            url=f"{cfg.FRONTEND_URL}/admin?token={token}",
            status_code=302,
        )

    except ImportError:
        raise HTTPException(503, "python3-saml not installed. Run: pip install python3-saml")


@router.get("/sso/{agency_slug}/metadata", tags=["SSO"])
async def sso_metadata(agency_slug: str, request: Request):
    """Return SP metadata XML — upload this to your IdP."""
    sso_cfg = _sso_configs.get(agency_slug)
    if not sso_cfg:
        raise HTTPException(404, "SSO not configured.")
    try:
        from onelogin.saml2.settings import OneLogin_Saml2_Settings
        settings = OneLogin_Saml2_Settings(_build_saml_settings(request, agency_slug, sso_cfg))
        metadata, errors = settings.get_sp_metadata()
        if errors:
            raise HTTPException(500, f"Metadata errors: {errors}")
        return Response(content=metadata, media_type="text/xml")
    except ImportError:
        return Response(
            content="<error>python3-saml not installed</error>",
            media_type="text/xml", status_code=503,
        )


# ── Helpers ───────────────────────────────────────────────────

def _first(lst):
    return lst[0] if lst else None


def _build_saml_settings(request: Request, slug: str, sso_cfg: dict) -> dict:
    base = str(request.base_url).rstrip("/")
    return {
        "strict": True,
        "debug": cfg.DEBUG if hasattr(cfg, "DEBUG") else False,
        "sp": {
            "entityId": f"{base}/sso/{slug}/metadata",
            "assertionConsumerService": {
                "url": f"{base}/sso/{slug}/acs",
                "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST",
            },
            "NameIDFormat": "urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress",
        },
        "idp": {
            "entityId": str(sso_cfg.get("idp_metadata_url", "")),
            "singleSignOnService": {
                "url": str(sso_cfg.get("idp_metadata_url", "")),
                "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect",
            },
        },
    }


def _build_saml_auth(request: Request, slug: str, sso_cfg: dict):
    from onelogin.saml2.auth import OneLogin_Saml2_Auth
    req = {
        "https": "on" if request.url.scheme == "https" else "off",
        "http_host": request.headers.get("host", "localhost"),
        "script_name": request.url.path,
        "get_data": dict(request.query_params),
        "post_data": {},
    }
    return OneLogin_Saml2_Auth(req, _build_saml_settings(request, slug, sso_cfg))