app.py

#!/usr/bin/env python3
"""Stateless SAML 2.0 Identity Provider — Flask application."""

import os
import secrets

from flask import Flask, make_response, redirect, render_template, request

from auth import (
    clear_session_cookie,
    create_jwt,
    get_current_user,
    set_session_cookie,
)
from config import Config
from saml import (
    SignatureVerificationError,
    build_metadata,
    build_saml_response,
    generate_self_signed_cert,
    parse_authn_request,
    verify_post_signature,
    verify_redirect_signature,
)
from scim import scim_bp
from sp_registry import get_sp
from user_store import UserStore

app = Flask(__name__)
app.register_blueprint(scim_bp)
store = UserStore()


def render_saml_response(user: dict, authn_req: dict, relay_state: str) -> str:
    """Build the signed SAML Response and wrap it in an auto-POST form."""
    saml_response = build_saml_response(
        user=user,
        request_id=authn_req["id"],
        acs_url=authn_req["acs_url"],
        sp_entity_id=authn_req["issuer"],
    )
    return render_template(
        "auto_post.html",
        acs_url=authn_req["acs_url"],
        saml_response=saml_response,
        relay_state=relay_state,
    )


@app.route("/sso", methods=["GET", "POST"])
def sso():
    """SSO endpoint — accepts SAMLRequest via GET (Redirect) or POST binding."""
    if request.method == "GET":
        saml_request = request.args.get("SAMLRequest", "")
        relay_state = request.args.get("RelayState", "")
        binding = "redirect"
    else:
        saml_request = request.form.get("SAMLRequest", "")
        relay_state = request.form.get("RelayState", "")
        binding = "post"

    if not saml_request:
        return "Missing SAMLRequest parameter", 400

    authn_req = parse_authn_request(saml_request, binding)

    # AuthnRequest signature verification (optional per config).
    # If required, lookup the SP in the registry and verify per-binding.
    if Config.REQUIRE_SIGNED_AUTHN_REQUESTS:
        try:
            _verify_authn_request_signature(authn_req, binding)
        except SignatureVerificationError as e:
            return f"AuthnRequest rejected: {e}", 401

    user = get_current_user()
    if user:
        return render_saml_response(user, authn_req, relay_state)

    return render_template(
        "login.html", saml_request=saml_request, relay_state=relay_state
    )


def _verify_authn_request_signature(authn_req: dict, binding: str) -> None:
    """Verify a signed AuthnRequest. Raises SignatureVerificationError on failure.

    Trust boundary: the SP must be pre-registered in sp_registry.json. An
    unregistered issuer is rejected outright — the IdP has no way to know
    whose cert to trust otherwise.
    """
    issuer = authn_req.get("issuer", "")
    sp = get_sp(issuer)
    if sp is None:
        raise SignatureVerificationError(
            f"issuer {issuer!r} not registered in sp_registry.json"
        )

    if binding == "redirect":
        # Redirect binding: signature is over the raw query string octets.
        # request.query_string is bytes; decode as ASCII (URL-encoded text).
        verify_redirect_signature(
            request.query_string.decode("ascii"),
            sp.cert_bytes(),
        )
    else:
        # POST binding: signature is enveloped in the XML itself. We need
        # the raw XML bytes — re-decode from the form SAMLRequest field.
        import base64

        raw = request.form.get("SAMLRequest", "")
        padded = raw + "=" * (-len(raw) % 4)
        xml_bytes = base64.b64decode(padded)
        verify_post_signature(xml_bytes, sp.cert_bytes())


@app.route("/sso/login", methods=["POST"])
def sso_login():
    """Handle login form submission."""
    email = request.form.get("email", "")
    password = request.form.get("password", "")
    saml_request = request.form.get("SAMLRequest", "")
    relay_state = request.form.get("RelayState", "")

    if not saml_request:
        return "Missing SAMLRequest", 400

    user = store.validate_credentials(email, password)
    if user is None:
        return (
            render_template(
                "login.html",
                saml_request=saml_request,
                relay_state=relay_state,
                error="Invalid credentials or account disabled",
            ),
            401,
        )

    user_dict = user.to_dict()

    # The SAMLRequest was originally sent via either binding; retry on failure.
    try:
        authn_req = parse_authn_request(saml_request, binding="redirect")
    except Exception:
        authn_req = parse_authn_request(saml_request, binding="post")

    token = create_jwt(user_dict)
    html = render_saml_response(user_dict, authn_req, relay_state)
    resp = make_response(html)
    set_session_cookie(resp, token)
    return resp


@app.route("/metadata", methods=["GET"])
def metadata():
    """Return IdP SAML metadata."""
    xml = build_metadata()
    resp = make_response(xml)
    resp.headers["Content-Type"] = "application/xml"
    return resp


@app.route("/slo", methods=["GET", "POST"])
def slo():
    """Single Logout endpoint."""
    resp = make_response(render_template("logged_out.html"))
    clear_session_cookie(resp)
    return resp


@app.route("/logout", methods=["GET"])
def logout():
    """Local logout — clear cookie and redirect to confirmation."""
    resp = make_response(redirect("/slo"))
    clear_session_cookie(resp)
    return resp


if __name__ == "__main__":
    generate_self_signed_cert()

    if not Config.JWT_SECRET:
        Config.JWT_SECRET = secrets.token_hex(32)
        print(
            "WARNING: JWT_SECRET not set — generated random secret. "
            "Sessions will be invalidated on every restart. "
            "Set JWT_SECRET in .env for stable sessions."
        )

    if not Config.SCIM_BEARER_TOKEN:
        print(
            "WARNING: SCIM_BEARER_TOKEN not set — SCIM endpoints are "
            "UNAUTHENTICATED. Anyone who can reach /scim/v2 can delete "
            "or deactivate users. Do not expose this service to untrusted "
            "networks without setting SCIM_BEARER_TOKEN."
        )

    Config.validate()

    print(f"IdP Entity ID: {Config.IDP_ENTITY_ID}")
    print(f"SSO URL:       {Config.IDP_SSO_URL}")
    print(f"Metadata:      {Config.IDP_ENTITY_ID}")

    port = int(os.environ.get("PORT", "5001"))
    app.run(host="0.0.0.0", port=port, debug=Config.FLASK_DEBUG)