Cortex-Inara/cortex/routers/onboarding.py

"""
Onboarding router — invite-based setup + persona creation.

Routes:
  GET  /setup/{token}      → show password setup form (step 1)
  POST /setup/{token}      → set password, redirect to persona step
  GET  /setup/persona      → show persona creation form (step 2, requires auth)
  POST /setup/persona      → create persona, redirect to /{user}/{persona}
"""

import logging
import re
from pathlib import Path

from fastapi import APIRouter, Form, Request
from fastapi.responses import HTMLResponse, RedirectResponse

from auth_utils import (
    COOKIE_NAME, validate_invite, consume_invite,
    set_password, create_token,
)
from persona_template import create_persona
from persona import list_user_personas, validate as validate_persona

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/setup")

_STATIC = Path(__file__).parent.parent / "static"
_SLUG_RE = re.compile(r"^[a-z_][a-z0-9_-]{0,31}$")


def _setup_page(error: str = "", step: int = 1) -> str:
    html = (_STATIC / "setup.html").read_text()
    if error:
        html = html.replace(
            "<!-- ERROR -->",
            f'<p class="error">{error}</p>',
        )
    if step == 2:
        html = html.replace("location.search)", "location.search)", 1)  # noop, handled by ?step=2
    return html


# ---------------------------------------------------------------------------
# Step 2 — persona creation (requires active session)
# IMPORTANT: must be registered before /{token} so "/persona" literal wins
# ---------------------------------------------------------------------------

@router.get("/persona", include_in_schema=False)
async def persona_page(request: Request):
    from auth_utils import decode_token
    import jwt
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return RedirectResponse("/login", status_code=302)
    try:
        decode_token(token)
    except jwt.InvalidTokenError:
        return RedirectResponse("/login", status_code=302)

    html = (_STATIC / "setup.html").read_text()
    # Show step 2 directly — inject ?step=2 behaviour inline
    html = html.replace(
        "if (params.get('step') === '2') {",
        "if (true || params.get('step') === '2') {",
    )
    return HTMLResponse(html)


@router.post("/persona", include_in_schema=False)
async def persona_submit(
    request: Request,
    step: str = Form(...),
    persona_name: str = Form(...),
    display_name: str = Form(...),
    user_real_name: str = Form(...),
    emoji: str = Form(default="✨"),
    description: str = Form(default=""),
):
    from auth_utils import decode_token
    import jwt

    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return RedirectResponse("/login", status_code=302)
    try:
        username = decode_token(token)
    except jwt.InvalidTokenError:
        return RedirectResponse("/login", status_code=302)

    # Validate persona slug
    if not _SLUG_RE.match(persona_name):
        html = (_STATIC / "setup.html").read_text().replace(
            "if (params.get('step') === '2') {",
            "if (true || params.get('step') === '2') {",
        ).replace("<!-- ERROR -->", '<p class="error">Invalid persona name. Use lowercase letters, digits, _ or - only.</p>')
        return HTMLResponse(html, status_code=422)

    # Check for collision
    existing = list_user_personas(username)
    if persona_name in existing:
        html = (_STATIC / "setup.html").read_text().replace(
            "if (params.get('step') === '2') {",
            "if (true || params.get('step') === '2') {",
        ).replace("<!-- ERROR -->", f'<p class="error">Persona "{persona_name}" already exists.</p>')
        return HTMLResponse(html, status_code=422)

    create_persona(
        username=username,
        persona_name=persona_name,
        display_name=display_name.strip() or persona_name.capitalize(),
        user_real_name=user_real_name.strip() or username.capitalize(),
        emoji=emoji or "✨",
        description=description.strip(),
    )
    logger.info("persona created: %s/%s", username, persona_name)
    return RedirectResponse(f"/{username}/{persona_name}", status_code=302)


# ---------------------------------------------------------------------------
# Step 1 — invite token → set password
# IMPORTANT: registered after /persona so the literal path wins above
# ---------------------------------------------------------------------------

@router.get("/{token}", include_in_schema=False)
async def setup_page(token: str, request: Request):
    """Show the password setup page for a valid invite token."""
    username = validate_invite(token)
    if not username:
        return HTMLResponse(
            "<h1 style='font-family:sans-serif;padding:2rem'>This link is invalid or has expired.</h1>",
            status_code=400,
        )
    return HTMLResponse(_setup_page())


@router.get("/{token}/persona", include_in_schema=False)
async def setup_persona_via_token(token: str, request: Request):
    """After password setup, redirect to the generic /setup/persona page."""
    # Cookie is already set — just redirect. Token is consumed so this is safe.
    return RedirectResponse("/setup/persona", status_code=302)


@router.post("/{token}", include_in_schema=False)
async def setup_submit(
    token: str,
    step: str = Form(...),
    password: str = Form(default=""),
    confirm: str = Form(default=""),
):
    username = validate_invite(token)
    if not username:
        return HTMLResponse(
            "<h1 style='font-family:sans-serif;padding:2rem'>This link is invalid or has expired.</h1>",
            status_code=400,
        )

    if step == "password":
        if len(password) < 8:
            return HTMLResponse(_setup_page("Password must be at least 8 characters."))
        if password != confirm:
            return HTMLResponse(_setup_page("Passwords do not match."))

        set_password(username, password)
        consume_invite(username)
        logger.info("setup complete (password): %s", username)

        # Log them in and move to persona step
        resp = RedirectResponse(f"/setup/{token}/persona", status_code=302)
        resp.set_cookie(
            COOKIE_NAME,
            create_token(username),
            max_age=30 * 86400,
            httponly=True,
            samesite="lax",
            secure=False,
        )
        return resp

    return HTMLResponse(_setup_page("Unknown step."), status_code=400)