FabledSteward/steward/auth/routes.py

from __future__ import annotations
import bcrypt
from quart import Blueprint, render_template, request, redirect, url_for, session
from sqlalchemy import select, func
from steward.auth.middleware import login_user, logout_user, public_route, safe_next_url
from steward.models.users import User, UserRole

auth_bp = Blueprint("auth", __name__)


async def _provision_external_user(app, username: str, email: str, role_str: str):
    """Create or update a user from an external auth provider (OIDC/LDAP).

    External users get a randomised unusable password so local login is blocked.
    Their role is updated on every login to reflect current IdP group membership.
    """
    import secrets
    from steward.models.users import UserRole
    async with app.db_sessionmaker() as db:
        result = await db.execute(select(User).where(User.username == username))
        user = result.scalar_one_or_none()
        new_role = UserRole(role_str)
        async with db.begin():
            if user is None:
                unusable_hash = bcrypt.hashpw(secrets.token_bytes(32), bcrypt.gensalt()).decode()
                user = User(username=username, email=email or f"{username}@external",
                            password_hash=unusable_hash, role=new_role)
                db.add(user)
            else:
                if user.role != new_role:
                    user.role = new_role
    # Re-fetch committed user
    async with app.db_sessionmaker() as db:
        result = await db.execute(select(User).where(User.username == username))
        return result.scalar_one()


async def get_user_count(app) -> int:
    async with app.db_sessionmaker() as db:
        result = await db.execute(select(func.count()).select_from(User))
        return result.scalar_one()


async def get_user_by_username(app, username: str) -> User | None:
    async with app.db_sessionmaker() as db:
        result = await db.execute(select(User).where(User.username == username))
        return result.scalar_one_or_none()


@auth_bp.get("/login")
@public_route("the login page")
async def login():
    from quart import current_app
    if await get_user_count(current_app) == 0:
        return redirect(url_for("auth.setup"))
    oidc_cfg = current_app.config.get("OIDC", {})
    return await render_template("auth/login.html",
                                 oidc_enabled=oidc_cfg.get("enabled", False),
                                 next_url=safe_next_url(request.args.get("next")))


@auth_bp.post("/login")
@public_route("the login form")
async def login_post():
    from quart import current_app
    from steward.core.audit import log_audit
    form = await request.form
    username = form.get("username", "").strip()
    password_str = form.get("password", "")
    password = password_str.encode()
    dest = safe_next_url(form.get("next")) or url_for("dashboard.index")

    # Try local auth first
    user = await get_user_by_username(current_app, username)
    if user and user.is_active and bcrypt.checkpw(password, user.password_hash.encode()):
        login_user(user)
        await log_audit(current_app, user.id, user.username, "auth.login",
                        detail={"method": "local", "role": user.role.value})
        return redirect(dest)

    # Try LDAP fallback if enabled
    ldap_cfg = current_app.config.get("LDAP", {})
    if ldap_cfg.get("enabled"):
        from steward.auth.ldap_auth import ldap_authenticate
        ldap_info = await ldap_authenticate(ldap_cfg, username, password_str)
        if ldap_info:
            user = await _provision_external_user(
                current_app, ldap_info["username"], ldap_info["email"], ldap_info["role"]
            )
            login_user(user)
            await log_audit(current_app, user.id, user.username, "auth.login",
                            detail={"method": "ldap", "role": user.role.value})
            return redirect(dest)

    oidc_cfg = current_app.config.get("OIDC", {})
    return await render_template(
        "auth/login.html",
        error="Invalid credentials",
        oidc_enabled=oidc_cfg.get("enabled", False),
        next_url=safe_next_url(form.get("next")),
    ), 400


@auth_bp.get("/login/oidc")
@public_route("starts SSO login")
async def login_oidc():
    from quart import current_app
    from steward.auth.oidc import get_discovery, build_authorize_url, generate_state
    oidc_cfg = current_app.config.get("OIDC", {})
    if not oidc_cfg.get("enabled") or not oidc_cfg.get("discovery_url"):
        return redirect(url_for("auth.login"))
    try:
        doc = await get_discovery(oidc_cfg["discovery_url"])
    except Exception as exc:
        return await render_template("auth/login.html", error=f"OIDC discovery failed: {exc}",
                                     oidc_enabled=True), 502
    state = generate_state()
    nonce = generate_state()
    session["oidc_state"] = state
    session["oidc_nonce"] = nonce
    session["oidc_next"] = safe_next_url(request.args.get("next"))
    redirect_uri = url_for("auth.login_oidc_callback", _external=True)
    url = build_authorize_url(
        doc, oidc_cfg["client_id"], redirect_uri,
        oidc_cfg.get("scopes", "openid profile email"), state, nonce,
    )
    return redirect(url)


@auth_bp.get("/login/oidc/callback")
@public_route("SSO login callback")
async def login_oidc_callback():
    from quart import current_app
    from steward.auth.oidc import get_discovery, exchange_code, get_userinfo, map_role
    from steward.core.audit import log_audit

    oidc_cfg = current_app.config.get("OIDC", {})
    error = request.args.get("error")
    if error:
        desc = request.args.get("error_description", error)
        return await render_template("auth/login.html", error=f"SSO error: {desc}",
                                     oidc_enabled=True), 400

    code = request.args.get("code", "")
    state = request.args.get("state", "")
    if not code or state != session.pop("oidc_state", None):
        return await render_template("auth/login.html", error="Invalid SSO state — try again.",
                                     oidc_enabled=True), 400

    try:
        doc = await get_discovery(oidc_cfg["discovery_url"])
        redirect_uri = url_for("auth.login_oidc_callback", _external=True)
        tokens = await exchange_code(
            doc, oidc_cfg["client_id"], oidc_cfg["client_secret"], code, redirect_uri
        )
        userinfo = await get_userinfo(doc, tokens["access_token"])
    except Exception as exc:
        return await render_template("auth/login.html", error=f"SSO login failed: {exc}",
                                     oidc_enabled=True), 502

    username_claim = oidc_cfg.get("username_claim", "preferred_username")
    email_claim = oidc_cfg.get("email_claim", "email")
    groups_claim = oidc_cfg.get("groups_claim", "groups")

    username = userinfo.get(username_claim) or userinfo.get("sub", "")
    email = userinfo.get(email_claim, "") or f"{username}@oidc"
    groups = userinfo.get(groups_claim, [])
    if not isinstance(groups, list):
        groups = []
    role = map_role(groups, oidc_cfg.get("admin_group", ""), oidc_cfg.get("operator_group", ""))

    if not username:
        return await render_template("auth/login.html",
                                     error="SSO returned no username.",
                                     oidc_enabled=True), 400

    user = await _provision_external_user(current_app, username, email, role)
    login_user(user)
    await log_audit(current_app, user.id, user.username, "auth.login",
                    detail={"method": "oidc", "role": role})
    dest = safe_next_url(session.pop("oidc_next", None)) or url_for("dashboard.index")
    return redirect(dest)


@auth_bp.get("/logout")
@public_route("ends the session")
async def logout():
    logout_user()
    return redirect(url_for("auth.login"))


@auth_bp.get("/setup")
@public_route("first-run setup (refuses once a user exists)")
async def setup():
    from quart import current_app
    if await get_user_count(current_app) > 0:
        return redirect(url_for("auth.login"))
    # Pre-fill the public URL with however the operator reached this page — the
    # common case is correct, and it's what agent installs / share links use.
    return await render_template(
        "auth/setup.html", default_url=request.host_url.rstrip("/"))


@auth_bp.post("/setup")
@public_route("first-run setup (refuses once a user exists)")
async def setup_post():
    from quart import current_app
    from steward.core.audit import log_audit
    if await get_user_count(current_app) > 0:
        return redirect(url_for("auth.login"))
    form = await request.form
    username = form.get("username", "").strip()
    email = form.get("email", "").strip()
    password = form.get("password", "").encode()
    if not username or not email or not password:
        return await render_template("auth/setup.html", error="All fields required"), 400
    pw_hash = bcrypt.hashpw(password, bcrypt.gensalt()).decode()
    user = User(username=username, email=email, password_hash=pw_hash, role=UserRole.admin)
    steward_url = (form.get("public_base_url", "") or "").strip().rstrip("/")
    async with current_app.db_sessionmaker() as db:
        async with db.begin():
            db.add(user)
            if steward_url:
                from steward.core.settings import set_setting
                await set_setting(db, "general.public_base_url", steward_url)
    if steward_url:
        # Apply immediately so agent installs / share links use it without a restart.
        current_app.config["PUBLIC_BASE_URL"] = steward_url
    login_user(user)
    await log_audit(current_app, user.id, user.username, "auth.setup",
                    detail={"username": username})
    return redirect(url_for("dashboard.index"))