FabledSteward/steward/auth/routes.py

from __future__ import annotations
from datetime import datetime, timezone

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.auth.throttle import login_throttle
from steward.core.proxy import client_ip
from steward.models.users import User, UserRole

auth_bp = Blueprint("auth", __name__)


_SOURCE_LABEL = {"ldap": "LDAP", "oidc": "SSO"}


async def _provision_external_user(app, source: str, subject: str, username: str,
                                   email: str, role_str: str) -> tuple[User | None, str | None]:
    """The Steward user for an identity a provider vouched for: (user, None),
    or (None, the reason to show at the login page).

    Found by (source, subject) only. A first sign-in creates the user with an
    unusable password; a username or email another account already holds is
    refused rather than merged, so a directory entry named `admin` cannot sign
    in as the local admin. The role follows the provider's groups on every
    sign-in; a user disabled in Steward stays out.
    """
    import secrets
    from sqlalchemy.exc import IntegrityError
    label = _SOURCE_LABEL[source]
    if not subject or len(subject) > 512 or not username or len(username) > 64:
        return None, f"{label} returned a username or id Steward cannot store (at most 64 characters)."
    new_role = UserRole(role_str)
    async with app.db_sessionmaker() as db:
        try:
            async with db.begin():
                user = (await db.execute(select(User).where(
                    User.auth_source == source, User.external_id == subject))).scalar_one_or_none()
                if user is not None:
                    if not user.is_active:
                        return None, "This account is disabled in Steward."
                    user.role = new_role
                else:
                    email = email or f"{username}@{source}"
                    clash = (await db.execute(select(User.username).where(
                        (User.username == username) | (User.email == email)))).first()
                    if clash is not None:
                        return None, (
                            f"Steward already has an account named {clash[0]} with this username "
                            f"or email, and it does not sign in with {label}. An admin must rename "
                            f"or remove it before this {label} account can sign in.")
                    user = User(username=username, email=email, role=new_role,
                                password_hash=bcrypt.hashpw(secrets.token_bytes(32),
                                                            bcrypt.gensalt()).decode(),
                                auth_source=source, external_id=subject)
                    db.add(user)
        except IntegrityError:
            # Another sign-in created the same name between our check and insert.
            return None, f"Could not create the {label} account; try signing in again."
        db.expunge(user)
    return user, None


async def _login_refused(message: str, next_url: str | None = None):
    """The login page with why a provider-vouched sign-in was still refused."""
    from quart import current_app
    return await render_template(
        "auth/login.html", error=message,
        oidc_enabled=current_app.config.get("OIDC", {}).get("enabled", False),
        next_url=safe_next_url(next_url)), 403


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")

    # Too many recent failures from this address: refuse before checking anything.
    ip = client_ip()
    wait_s = login_throttle.retry_after(
        ip, current_app.config.get("LOGIN_MAX_FAILURES", 5),
        current_app.config.get("LOGIN_LOCKOUT_MINUTES", 15) * 60)
    if wait_s:
        oidc_cfg = current_app.config.get("OIDC", {})
        return await render_template(
            "auth/login.html",
            error=f"Too many failed logins from this address. Try again in "
                  f"{max(1, round(wait_s / 60))} minute{'s' if wait_s > 90 else ''}.",
            oidc_enabled=oidc_cfg.get("enabled", False),
            next_url=safe_next_url(form.get("next")),
        ), 429

    # Try local auth first
    user = await get_user_by_username(current_app, username)
    if (user and user.auth_source == "local" and user.is_active
            and bcrypt.checkpw(password, user.password_hash.encode())):
        login_throttle.succeeded(ip)
        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, refusal = await _provision_external_user(
                current_app, "ldap", ldap_info["subject"], ldap_info["username"],
                ldap_info["email"], ldap_info["role"])
            if refusal:
                return await _login_refused(refusal, form.get("next"))
            login_throttle.succeeded(ip)
            login_user(user)
            await log_audit(current_app, user.id, user.username, "auth.login",
                            detail={"method": "ldap", "role": user.role.value})
            return redirect(dest)

    login_throttle.failed(ip)
    await log_audit(current_app, None, username or "(blank)", "auth.login_failed",
                    detail={"address": ip})
    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,
                                    id_token_refusal)
    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

    nonce = session.pop("oidc_nonce", None)
    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

    refusal = id_token_refusal(tokens.get("id_token", ""), nonce, userinfo.get("sub", ""))
    if refusal:
        return await render_template("auth/login.html", error=refusal, oidc_enabled=True), 400

    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")

    subject = userinfo["sub"]
    username = userinfo.get(username_claim) or subject
    email = userinfo.get(email_claim, "")
    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", ""))

    user, refusal = await _provision_external_user(current_app, "oidc", subject, username, email, role)
    if refusal:
        return await _login_refused(refusal, session.pop("oidc_next", None))
    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.post("/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"))


# ── Invites (milestone 496) ──────────────────────────────────────────────────

_INVITE_GONE = ("This invite link has expired, been used or been withdrawn. "
                "Ask whoever sent it for a new one.")


async def _invite_page(key: str, invite=None, inviter: str = "", error: str | None = None,
                       form=None, status: int = 200):
    """The invite page: the form while `invite` works, else one plain message
    that doesn't say whether the key was unknown, used, revoked or expired."""
    if invite is None and error is None:
        error = _INVITE_GONE
    return await render_template(
        "auth/invite.html", key=key, invite=invite, inviter=inviter, error=error,
        form=form or {}, min_password=_min_password()), status


def _min_password() -> int:
    from steward.settings.users import MIN_PASSWORD
    return MIN_PASSWORD


async def _inviter_name(db, invite) -> str:
    if invite.invited_by is None:
        return "An admin"
    who = await db.get(User, invite.invited_by)
    return who.username if who else "An admin"


def _invite_throttled() -> str | None:
    """Why this address may not try an invite key now (too many bad ones), or None."""
    from quart import current_app
    wait_s = login_throttle.retry_after(
        client_ip(), current_app.config.get("LOGIN_MAX_FAILURES", 5),
        current_app.config.get("LOGIN_LOCKOUT_MINUTES", 15) * 60)
    if not wait_s:
        return None
    minutes = max(1, round(wait_s / 60))
    return f"Too many tries from this address. Try again in {minutes} minute{'s' if minutes != 1 else ''}."


@auth_bp.get("/invite/<key>")
@public_route("an invite link: the key is the credential")
async def invite(key: str):
    from quart import current_app
    from steward.core.invites import find_pending
    if wait := _invite_throttled():
        return await _invite_page(key, error=wait, status=429)
    async with current_app.db_sessionmaker() as db:
        found = await find_pending(db, key)
        inviter = await _inviter_name(db, found) if found else ""
    if found is None:
        login_throttle.failed(client_ip())
        return await _invite_page(key, status=404)
    return await _invite_page(key, found, inviter, form={"email": found.email})


@auth_bp.post("/invite/<key>")
@public_route("an invite link: the key is the credential")
async def invite_accept(key: str):
    """Make the invitee's local account from their own username, email and
    password, with the role the admin chose; use the key up and sign them in,
    in one transaction, so a key makes one account."""
    from quart import current_app
    from sqlalchemy.exc import IntegrityError
    from steward.core.audit import log_audit
    from steward.core.invites import find_pending, signup_problem, taken
    from steward.settings.users import _hash
    if wait := _invite_throttled():
        return await _invite_page(key, error=wait, status=429)
    form = await request.form
    username = form.get("username") or ""
    email = (form.get("email") or "").strip()
    password, confirm = form.get("password") or "", form.get("confirm") or ""
    shown = {"username": username, "email": email}
    try:
        async with current_app.db_sessionmaker() as db:
            async with db.begin():
                found = await find_pending(db, key, lock=True)
                if found is None:
                    login_throttle.failed(client_ip())
                    return await _invite_page(key, status=404)
                inviter = await _inviter_name(db, found)
                problem = signup_problem(username, email, password, confirm) or await taken(db, username, email)
                if problem:
                    return await _invite_page(key, found, inviter, problem, shown, 400)
                user = User(username=username, email=email, role=found.role, auth_source="local",
                            password_hash=_hash(password))
                db.add(user)
                await db.flush()
                found.used_at = datetime.now(timezone.utc)
                found.used_by = user.id
                invite_id, role = found.id, found.role.value
            db.expunge(user)
    except IntegrityError:
        async with current_app.db_sessionmaker() as db:
            found = await find_pending(db, key)
            inviter = await _inviter_name(db, found) if found else ""
        if found is None:
            return await _invite_page(key, status=404)
        return await _invite_page(key, found, inviter,
                                  "That username or email was just taken. Pick another.", shown, 400)
    login_user(user)
    await log_audit(current_app, user.id, user.username, "user.joined", "user", user.id,
                    detail={"username": user.username, "role": role, "invite": invite_id,
                            "invited_by": inviter})
    return redirect(url_for("dashboard.index"))