FabledScribe/src/scribe/routes/auth.py

import asyncio
import base64
import hashlib
import os
import secrets

from quart import Blueprint, g, jsonify, redirect, request, session

from scribe.auth import login_required, get_current_user_id
from scribe.config import Config
from scribe.rate_limit import is_rate_limited
from scribe.services.auth import (
    authenticate,
    change_password,
    create_password_reset_token,
    create_user,
    get_user_by_email,
    get_user_by_id,
    get_user_by_username,
    get_user_count,
    invalidate_other_sessions,
    is_registration_open,
    register_with_invitation,
    reset_password_with_token,
    update_user_email,
    validate_invitation_token,
    verify_password,
)
from scribe.services.logging import log_audit
from scribe.services.notifications import (
    notify_security_event,
    send_password_reset_email,
    send_password_reset_success_email,
)
from scribe.services.email import get_base_url, is_smtp_configured
from scribe.services.oauth import (
    build_auth_url,
    exchange_code,
    get_userinfo,
    find_or_create_oauth_user,
)

auth_bp = Blueprint("auth", __name__, url_prefix="/api/auth")


def _client_ip() -> str:
    if Config.TRUST_PROXY_HEADERS:
        for header in ("X-Forwarded-For", "X-Real-IP"):
            val = request.headers.get(header, "").split(",")[0].strip()
            if val:
                return val
    return request.remote_addr or ""


@auth_bp.route("/register", methods=["POST"])
async def register():
    if not Config.LOCAL_AUTH_ENABLED:
        return jsonify({"error": "Local authentication is disabled. Please use SSO."}), 403
    if await is_rate_limited(f"register:{_client_ip()}", max_requests=5, window_seconds=300):
        return jsonify({"error": "Too many registration attempts. Try again later."}), 429
    if not await is_registration_open():
        return jsonify({"error": "Registration is closed"}), 403

    data = await request.get_json()
    username = (data.get("username") or "").strip()
    password = data.get("password") or ""
    email = (data.get("email") or "").strip() or None

    if not username:
        return jsonify({"error": "Username is required"}), 400
    if len(password) < 8:
        return jsonify({"error": "Password must be at least 8 characters"}), 400

    try:
        user = await create_user(username, password, email)
    except Exception:
        return jsonify({"error": "Username already taken"}), 409

    session["user_id"] = user.id
    session["session_version"] = user.session_version
    await log_audit("register", user_id=user.id, username=user.username, ip_address=_client_ip())
    return jsonify(user.to_dict()), 201


@auth_bp.route("/login", methods=["POST"])
async def login():
    if not Config.LOCAL_AUTH_ENABLED:
        return jsonify({"error": "Local authentication is disabled. Please use SSO."}), 403
    if await is_rate_limited(f"login:{_client_ip()}", max_requests=10, window_seconds=60):
        return jsonify({"error": "Too many login attempts. Try again later."}), 429
    data = await request.get_json()
    username = (data.get("username") or "").strip()
    password = data.get("password") or ""

    if not username or not password:
        return jsonify({"error": "Username and password are required"}), 400

    user = await authenticate(username, password)
    if not user:
        await log_audit("login_failed", username=username, ip_address=_client_ip())
        # Try to notify the actual user about the failed attempt
        target_user = await get_user_by_username(username)
        if target_user:
            asyncio.create_task(notify_security_event(
                target_user.id, "login_failed",
                {"ip_address": _client_ip(), "username": username},
            ))
        return jsonify({"error": "Invalid username or password"}), 401

    session["user_id"] = user.id
    session["session_version"] = user.session_version
    await log_audit("login", user_id=user.id, username=user.username, ip_address=_client_ip())
    asyncio.create_task(notify_security_event(
        user.id, "login",
        {"ip_address": _client_ip()},
    ))
    return jsonify(user.to_dict())


@auth_bp.route("/logout", methods=["POST"])
async def logout():
    uid = session.get("user_id")
    if uid:
        user = await get_user_by_id(uid)
        await log_audit("logout", user_id=uid, username=user.username if user else None, ip_address=_client_ip())
        asyncio.create_task(notify_security_event(
            uid, "logout",
            {"ip_address": _client_ip()},
        ))
    session.clear()
    return jsonify({"status": "ok"})


@auth_bp.route("/me", methods=["GET"])
@login_required
async def me():
    user = await get_user_by_id(get_current_user_id())
    if not user:
        return jsonify({"error": "User not found"}), 404
    return jsonify(user.to_dict())


@auth_bp.route("/password", methods=["PUT"])
@login_required
async def update_password():
    data = await request.get_json()
    current = data.get("current_password") or ""
    new_pw = data.get("new_password") or ""

    if not current:
        return jsonify({"error": "Current password is required"}), 400
    if len(new_pw) < 8:
        return jsonify({"error": "New password must be at least 8 characters"}), 400

    uid = get_current_user_id()
    new_version = await change_password(uid, current, new_pw)
    if new_version is None:
        return jsonify({"error": "Current password is incorrect"}), 403

    # Keep the current session alive with the new version (other sessions are now invalidated)
    session["session_version"] = new_version
    await log_audit("password_change", user_id=uid, username=g.user.username, ip_address=_client_ip())
    asyncio.create_task(notify_security_event(
        uid, "password_change",
        {"ip_address": _client_ip()},
    ))
    return jsonify({"status": "ok"})


@auth_bp.route("/invalidate-sessions", methods=["POST"])
@login_required
async def invalidate_sessions_route():
    """Invalidate all other active sessions by bumping the session version.

    The current session is kept alive. Useful after an SSO/OAuth password change
    where the app has no visibility into credential rotation.
    """
    uid = get_current_user_id()
    new_version = await invalidate_other_sessions(uid)
    session["session_version"] = new_version
    await log_audit("sessions_invalidated", user_id=uid, username=g.user.username, ip_address=_client_ip())
    return jsonify({"status": "ok"})


@auth_bp.route("/email", methods=["PUT"])
@login_required
async def update_email():
    data = await request.get_json()
    new_email = (data.get("email") or "").strip().lower() or None
    password = data.get("password") or ""

    uid = get_current_user_id()
    user = await get_user_by_id(uid)
    if not user:
        return jsonify({"error": "User not found"}), 404

    # Require password confirmation only for local-auth users
    if user.password_hash is not None:
        if not password:
            return jsonify({"error": "Password is required to change your email"}), 400
        if not verify_password(password, user.password_hash):
            return jsonify({"error": "Incorrect password"}), 403

    # Check the new email isn't taken by a different account
    if new_email:
        existing = await get_user_by_email(new_email)
        if existing and existing.id != uid:
            return jsonify({"error": "Email already in use by another account"}), 409

    await update_user_email(uid, new_email)
    await log_audit("email_change", user_id=uid, username=user.username, ip_address=_client_ip())
    updated = await get_user_by_id(uid)
    return jsonify(updated.to_dict())


@auth_bp.route("/forgot-password", methods=["POST"])
async def forgot_password():
    if await is_rate_limited(f"forgot:{_client_ip()}", max_requests=5, window_seconds=300):
        return jsonify({"status": "ok"})
    data = await request.get_json()
    email = (data.get("email") or "").strip().lower()

    if not email:
        return jsonify({"error": "Email is required"}), 400

    # Always return success to prevent email enumeration
    user = await get_user_by_email(email)
    if user and await is_smtp_configured():
        raw_token = await create_password_reset_token(user.id)
        base_url = await get_base_url()
        reset_url = f"{base_url}/reset-password?token={raw_token}"
        asyncio.create_task(send_password_reset_email(user.email, reset_url))
        await log_audit(
            "password_reset_requested",
            user_id=user.id,
            username=user.username,
            ip_address=_client_ip(),
        )

    return jsonify({"status": "ok"})


@auth_bp.route("/reset-password", methods=["POST"])
async def reset_password():
    if await is_rate_limited(f"reset:{_client_ip()}", max_requests=10, window_seconds=60):
        return jsonify({"error": "Too many attempts. Try again later."}), 429
    data = await request.get_json()
    token = (data.get("token") or "").strip()
    new_password = data.get("new_password") or ""

    if not token:
        return jsonify({"error": "Reset token is required"}), 400
    if len(new_password) < 8:
        return jsonify({"error": "Password must be at least 8 characters"}), 400

    user_id = await reset_password_with_token(token, new_password)
    if user_id is None:
        return jsonify({"error": "Invalid or expired reset link"}), 400

    user = await get_user_by_id(user_id)
    if user:
        await log_audit(
            "password_reset_completed",
            user_id=user.id, username=user.username, ip_address=_client_ip(),
        )
        if user.email:
            asyncio.create_task(send_password_reset_success_email(user.email))
    return jsonify({"status": "ok"})


@auth_bp.route("/invitation/<token>", methods=["GET"])
async def check_invitation(token: str):
    invitation = await validate_invitation_token(token)
    if not invitation:
        return jsonify({"valid": False})
    return jsonify({"valid": True, "email": invitation.email})


@auth_bp.route("/register-with-invite", methods=["POST"])
async def register_with_invite():
    data = await request.get_json()
    raw_token = (data.get("token") or "").strip()
    username = (data.get("username") or "").strip()
    password = data.get("password") or ""

    if not raw_token:
        return jsonify({"error": "Invitation token is required"}), 400
    if not username:
        return jsonify({"error": "Username is required"}), 400
    if len(password) < 8:
        return jsonify({"error": "Password must be at least 8 characters"}), 400

    try:
        user = await register_with_invitation(raw_token, username, password)
    except Exception:
        return jsonify({"error": "Username already taken"}), 409

    if not user:
        return jsonify({"error": "Invalid or expired invitation"}), 400

    session["user_id"] = user.id
    session["session_version"] = user.session_version
    await log_audit(
        "register_with_invite",
        user_id=user.id,
        username=user.username,
        ip_address=_client_ip(),
    )
    return jsonify(user.to_dict()), 201


@auth_bp.route("/status", methods=["GET"])
async def status():
    count = await get_user_count()
    reg_open = await is_registration_open()
    return jsonify({
        "has_users": count > 0,
        "registration_open": reg_open,
        "oauth_enabled": Config.oidc_enabled(),
        "local_auth_enabled": Config.LOCAL_AUTH_ENABLED,
    })


@auth_bp.route("/oauth/login", methods=["GET"])
async def oauth_login():
    if not Config.oidc_enabled():
        return jsonify({"error": "SSO is not configured"}), 404

    state = secrets.token_hex(32)
    code_verifier = secrets.token_urlsafe(64)

    digest = hashlib.sha256(code_verifier.encode()).digest()
    code_challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()

    session["oauth_state"] = state
    session["oauth_code_verifier"] = code_verifier

    url = await build_auth_url(state, code_challenge)
    return redirect(url)


@auth_bp.route("/oauth/callback", methods=["GET"])
async def oauth_callback():
    error = request.args.get("error")
    code = request.args.get("code")
    state = request.args.get("state")

    stored_state = session.pop("oauth_state", None)
    code_verifier = session.pop("oauth_code_verifier", None)

    if error or not code or state != stored_state:
        return redirect("/login?error=oauth")

    redirect_uri = Config.BASE_URL.rstrip("/") + "/api/auth/oauth/callback"
    try:
        token_response = await exchange_code(code, redirect_uri, code_verifier)
        claims = await get_userinfo(token_response["access_token"])
    except Exception:
        return redirect("/login?error=oauth")

    sub = claims.get("sub", "")
    # Only trust the email claim for account linking if the provider has verified it.
    # An unverified email could be used to hijack an existing local account.
    email_verified = claims.get("email_verified", False)
    email = claims.get("email", "") if email_verified else ""
    preferred_username = claims.get("preferred_username", "")

    if not sub:
        return redirect("/login?error=oauth")

    user = await find_or_create_oauth_user(sub, email, preferred_username)
    session["user_id"] = user.id
    session["session_version"] = user.session_version
    await log_audit("oauth_login", user_id=user.id, username=user.username, ip_address=_client_ip())
    return redirect("/")