pagepiper/app/cloud_session.py

# app/cloud_session.py
"""Cloud session auth for Pagepiper — validates cf_session cookie via Directus + Heimdall.

In local mode (CLOUD_MODE unset or false), require_paid_tier is a no-op.
In cloud mode, the Caddy proxy forwards the browser's Cookie header as
X-CF-Session. This module extracts cf_session, validates it against
Directus /users/me, then checks the user's Pagepiper tier via Heimdall.
Auto-provisions a free tier key for new users.
"""
from __future__ import annotations

import logging
import os
import re

import httpx
from fastapi import HTTPException, Request

log = logging.getLogger(__name__)

CLOUD_MODE: bool = os.environ.get("CLOUD_MODE", "").lower() in ("1", "true", "yes")
DIRECTUS_URL: str = os.environ.get("DIRECTUS_URL", "http://172.31.0.3:8055").rstrip("/")
HEIMDALL_URL: str = os.environ.get("HEIMDALL_URL", "https://license.circuitforge.tech").rstrip("/")
HEIMDALL_ADMIN_TOKEN: str = os.environ.get("HEIMDALL_ADMIN_TOKEN", "")

_TIER_ORDER = {"free": 0, "paid": 1, "premium": 2, "ultra": 3}


def _extract_session_token(cookie_header: str) -> str:
    m = re.search(r"(?:^|;)\s*cf_session=([^;]+)", cookie_header)
    return m.group(1).strip() if m else ""


def _get_user_id(jwt: str) -> str | None:
    try:
        resp = httpx.get(
            f"{DIRECTUS_URL}/users/me",
            headers={"Authorization": f"Bearer {jwt}"},
            timeout=5.0,
        )
        if resp.status_code == 200:
            return resp.json().get("data", {}).get("id")
    except Exception as exc:
        log.warning("Directus session check failed: %s", exc)
    return None


def _ensure_provisioned(user_id: str) -> None:
    if not HEIMDALL_ADMIN_TOKEN:
        return
    try:
        httpx.post(
            f"{HEIMDALL_URL}/admin/provision",
            json={"directus_user_id": user_id, "product": "pagepiper", "tier": "free"},
            headers={"Authorization": f"Bearer {HEIMDALL_ADMIN_TOKEN}"},
            timeout=5.0,
        )
    except Exception as exc:
        log.warning("Heimdall provision failed for user %s: %s", user_id, exc)


def _get_tier(user_id: str) -> str:
    if not HEIMDALL_ADMIN_TOKEN:
        return "free"
    try:
        resp = httpx.get(
            f"{HEIMDALL_URL}/admin/cloud/resolve",
            params={"directus_user_id": user_id, "product": "pagepiper"},
            headers={"Authorization": f"Bearer {HEIMDALL_ADMIN_TOKEN}"},
            timeout=5.0,
        )
        if resp.status_code == 200:
            return resp.json().get("tier", "free")
    except Exception as exc:
        log.warning("Heimdall tier check failed for user %s: %s", user_id, exc)
    return "free"


def resolve_authenticated_user(request: Request) -> str:
    """Validate the session cookie and return the Directus user_id. Raises 401 if invalid."""
    cookie_header = request.headers.get("x-cf-session", "")
    jwt = _extract_session_token(cookie_header)

    if not jwt:
        raise HTTPException(
            status_code=401,
            detail={
                "error": "auth_required",
                "message": "Sign in at circuitforge.tech to use Pagepiper cloud.",
            },
        )

    user_id = _get_user_id(jwt)
    if not user_id:
        raise HTTPException(
            status_code=401,
            detail={
                "error": "session_invalid",
                "message": "Your session has expired. Sign in again at circuitforge.tech.",
            },
        )

    _ensure_provisioned(user_id)
    return user_id


def require_paid_tier(request: Request) -> str:
    """FastAPI dependency — 401 if no valid session, 402 if tier < paid. Returns user_id.

    In local mode (CLOUD_MODE not set), returns LOCAL_USER_ID without any auth check.
    """
    if not CLOUD_MODE:
        from app.config import LOCAL_USER_ID
        return LOCAL_USER_ID

    user_id = resolve_authenticated_user(request)
    tier = _get_tier(user_id)

    if _TIER_ORDER.get(tier, 0) < _TIER_ORDER["paid"]:
        raise HTTPException(
            status_code=402,
            detail={
                "error": "upgrade_required",
                "message": (
                    "RAG chat requires a Paid tier Pagepiper license. "
                    "Upgrade at circuitforge.tech/software/pagepiper."
                ),
            },
        )

    return user_id