Cortex-Inara/cortex/routers/ui.py

"""
UI router — serves the web interface and handles login/logout.

Routes:
  GET  /              → redirect to /{user}/{persona} if logged in, else /login
  GET  /login         → login page
  POST /login         → validate credentials, set cookie, redirect
  POST /logout        → clear cookie, redirect to /login
  GET  /{user}/{persona}  → serve index.html with CORTEX_CONFIG injected
  GET  /{user}/{persona}/  → same (trailing slash)
"""

import logging
from pathlib import Path

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

from auth_utils import COOKIE_NAME, check_credentials, create_token, decode_token
from persona import list_user_personas, validate as validate_persona

logger = logging.getLogger(__name__)
router = APIRouter()

_STATIC = Path(__file__).parent.parent / "static"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _get_session_user(request: Request) -> str | None:
    """Return the authenticated username from the session cookie, or None."""
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return None
    try:
        return decode_token(token)
    except jwt.InvalidTokenError:
        return None


def _set_cookie(response: Response, username: str) -> None:
    from auth_utils import create_token
    from config import settings
    token = create_token(username)
    response.set_cookie(
        COOKIE_NAME,
        token,
        max_age=settings.jwt_expire_days * 86400,
        httponly=True,
        samesite="lax",
        secure=False,   # set True in production behind HTTPS
    )


def _first_persona(username: str) -> str | None:
    """Return the first available persona for a user, or None."""
    names = list_user_personas(username)
    return names[0] if names else None


# ---------------------------------------------------------------------------
# Root redirect
# ---------------------------------------------------------------------------

@router.get("/", include_in_schema=False)
async def root(request: Request):
    user = _get_session_user(request)
    if not user:
        return RedirectResponse("/login", status_code=302)
    persona = _first_persona(user)
    if not persona:
        return HTMLResponse("<h1>No personas configured for your account.</h1>", status_code=500)
    return RedirectResponse(f"/{user}/{persona}", status_code=302)


# ---------------------------------------------------------------------------
# Login / logout
# ---------------------------------------------------------------------------

@router.get("/login", include_in_schema=False)
async def login_page(request: Request):
    user = _get_session_user(request)
    if user:
        # Already logged in — redirect home
        persona = _first_persona(user)
        if persona:
            return RedirectResponse(f"/{user}/{persona}", status_code=302)
    return HTMLResponse((_STATIC / "login.html").read_text())


@router.post("/login", include_in_schema=False)
async def login(
    request: Request,
    username: str = Form(...),
    password: str = Form(...),
):
    if not check_credentials(username, password):
        logger.warning("failed login attempt for user: %s", username)
        html = (_STATIC / "login.html").read_text().replace(
            "<!-- ERROR -->",
            '<p class="error">Invalid username or password.</p>',
        )
        return HTMLResponse(html, status_code=401)

    persona = _first_persona(username)
    if not persona:
        return HTMLResponse("<h1>No personas configured for your account.</h1>", status_code=500)

    logger.info("login: %s", username)
    resp = RedirectResponse(f"/{username}/{persona}", status_code=302)
    _set_cookie(resp, username)
    return resp


@router.post("/logout", include_in_schema=False)
async def logout():
    resp = RedirectResponse("/login", status_code=302)
    resp.delete_cookie(COOKIE_NAME)
    return resp


# ---------------------------------------------------------------------------
# Main UI — /{username}/{persona}
# ---------------------------------------------------------------------------

@router.get("/{username}/{persona}", include_in_schema=False)
@router.get("/{username}/{persona}/", include_in_schema=False)
async def serve_ui(username: str, persona: str, request: Request):
    # Auth check
    session_user = _get_session_user(request)
    if not session_user:
        return RedirectResponse("/login", status_code=302)
    if session_user != username:
        return RedirectResponse(f"/{session_user}/{_first_persona(session_user) or ''}", status_code=302)

    # Validate persona exists
    try:
        validate_persona(username, persona)
    except ValueError:
        return RedirectResponse(f"/{username}/{_first_persona(username) or ''}", status_code=302)

    # Serve index.html with user/persona injected
    html = (_STATIC / "index.html").read_text()
    config_tag = (
        f'<script>window.CORTEX_CONFIG = '
        f'{{user: "{username}", persona: "{persona}"}};</script>'
    )
    html = html.replace("</head>", f"{config_tag}\n</head>", 1)
    return HTMLResponse(html)