Blame for managed_auth.py

# -*- coding: utf-8 -*-
"""SmartPaper edge-first authentication and browser task/payment bridge.

Cloudflare owns the durable browser session, quota, tasks, and payments. The
Hugging Face process verifies the signed edge assertion but does not call the
Worker over the server-to-server network. Authenticated business operations are
executed by the user's browser through the same-origin invisible component.
"""
from __future__ import annotations

import base64
import hashlib
import hmac
import html
import io
import json
import os
import secrets
import time
from typing import Any, Callable
from urllib.parse import unquote, urlsplit

import qrcode
import streamlit as st

from i18n import tr
from smartpaper_restore_bridge import edge_bridge

EDGE_ASSERTION_COOKIE_NAME = "smartpaper_edge_assertion"
EDGE_ASSERTION_AUDIENCE = "smartpaper-streamlit"
ASSERTION_CLOCK_SKEW_SECONDS = 30
def _runtime_value(name: str) -> str:
    value = os.environ.get(name, "")
    if not value:
        try:
            value = st.secrets.get(name, "")
        except Exception:
            value = ""
    return str(value or "").strip()


def _api_base() -> str:
    value = str(st.session_state.get("sp_auth_api_origin") or "").strip().rstrip("/")
    if not value:
        value = _runtime_value("SMARTPAPER_AUTH_API").rstrip("/")
    parsed = urlsplit(value)
    if (
        parsed.scheme != "https"
        or not parsed.netloc
        or parsed.username is not None
        or parsed.password is not None
        or parsed.path not in ("", "/")
        or parsed.query
        or parsed.fragment
    ):
        raise RuntimeError("SMARTPAPER_AUTH_API must be an HTTPS origin.")
    return f"https://{parsed.netloc}"


def _service_key() -> str:
    value = _runtime_value("SMARTPAPER_SERVICE_KEY")
    if not value:
        raise RuntimeError("SMARTPAPER_SERVICE_KEY is not configured.")
    return value


def _context_cookie(name: str) -> str:
    try:
        cookies = st.context.cookies
        value = cookies.get(name, "")
        if value:
            return unquote(str(value).strip())
        target = name.lower()
        for key, item in cookies.items():
            if str(key).lower() == target:
                return unquote(str(item or "").strip())
    except Exception:
        return ""
    return ""


def _decode_b64url(value: str) -> bytes:
    return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4))


def _verify_edge_assertion(token: str) -> dict[str, Any]:
    try:
        encoded, supplied_signature = token.split(".", 1)
    except ValueError as exc:
        raise RuntimeError("invalid_edge_assertion") from exc
    expected_signature = hmac.new(
        _service_key().encode("utf-8"),
        encoded.encode("ascii"),
        hashlib.sha256,
    ).hexdigest()
    if not hmac.compare_digest(expected_signature, supplied_signature):
        raise RuntimeError("invalid_edge_assertion_signature")
    try:
        payload = json.loads(_decode_b64url(encoded).decode("utf-8"))
    except Exception as exc:
        raise RuntimeError("invalid_edge_assertion_payload") from exc
    now = int(time.time())
    if payload.get("ver") != 1 or payload.get("aud") != EDGE_ASSERTION_AUDIENCE:
        raise RuntimeError("invalid_edge_assertion_claims")
    api_origin = str(payload.get("api_origin") or "").strip().rstrip("/")
    parsed = urlsplit(api_origin)
    if (
        parsed.scheme != "https"
        or not parsed.netloc
        or parsed.username is not None
        or parsed.password is not None
        or parsed.path not in ("", "/")
        or parsed.query
        or parsed.fragment
    ):
        raise RuntimeError("invalid_edge_assertion_origin")
    issued_at = int(payload.get("iat") or 0)
    expires_at = int(payload.get("exp") or 0)
    if issued_at > now + ASSERTION_CLOCK_SKEW_SECONDS or expires_at <= now - ASSERTION_CLOCK_SKEW_SECONDS:
        raise RuntimeError("expired_edge_assertion")
    if expires_at - issued_at > 9 * 60 * 60:
        raise RuntimeError("invalid_edge_assertion_lifetime")
    user = payload.get("user") or {}
    if not isinstance(user, dict) or not user.get("id") or user.get("status") != "active":
        raise RuntimeError("invalid_edge_assertion_identity")
    payload["api_origin"] = f"https://{parsed.netloc}"
    return payload


def _clear_generation_session_state() -> None:
    for key in (
        "sp_generation_task", "sp_tracked_task_id", "sp_current_node",
        "sp_task_recovery_initialized", "sp_last_ui_task_status",
        "sp_generation_idempotency_key", "zip_buffer", "zip_filename",
        "is_running", "start_pipeline", "resume_pipeline", "run_mode",
        "rag_engine", "up_amj_ref", "up_amj_data", "up_asq_ref",
        "up_asq_data", "sp_payment_required", "sp_payment_order",
        "sp_payment_request_id",
    ):
        st.session_state.pop(key, None)


def _clear_session() -> None:
    for key in (
        "sp_session_expires_at", "sp_user",
        "sp_last_verified_at", "sp_edge_assertion", "sp_auth_api_origin",
        "sp_product_ui_projection", "sp_product_tier",
    ):
        st.session_state.pop(key, None)


def require_alipay_login(language: str = "zh") -> dict[str, Any]:
    """Return the user authenticated by the Cloudflare edge gateway."""
    token = _context_cookie(EDGE_ASSERTION_COOKIE_NAME)
    if token and token == st.session_state.get("sp_edge_assertion") and st.session_state.get("sp_user"):
        return dict(st.session_state["sp_user"])
    if not token:
        _clear_session()
        st.stop()
    try:
        payload = _verify_edge_assertion(token)
    except Exception:
        _clear_session()
        st.stop()
    user = dict(payload["user"])
    st.session_state["sp_edge_assertion"] = token
    st.session_state["sp_session_expires_at"] = payload["exp"]
    st.session_state["sp_user"] = user
    st.session_state["sp_auth_api_origin"] = str(payload["api_origin"]).rstrip("/")
    st.session_state["sp_last_verified_at"] = time.time()
    return user


def _bridge_key(prefix: str, value: str) -> str:
    digest = hashlib.sha256(str(value).encode("utf-8")).hexdigest()[:20]
    return f"{prefix}_{digest}"


def _edge_call(
    action: str,
    command_id: str,
    payload: dict[str, Any] | None = None,
    *,
    key_prefix: str = "sp_edge",
) -> dict[str, Any] | None:
    return edge_bridge(
        action=action,
        command_id=command_id,
        api_origin=_api_base(),
        payload=payload or {},
        key=_bridge_key(key_prefix, command_id),
    )


def _raise_edge_error(result: dict[str, Any], fallback: str) -> None:
    message = (
        result.get("message")
        or result.get("reason")
        or result.get("error")
        or fallback
    )
    raise RuntimeError(str(message))


def _merge_quota_into_session(result: dict[str, Any]) -> None:
    """Refresh the approved account card immediately after a quota mutation."""
    quota = result.get("quota")
    if not isinstance(quota, dict):
        return
    user = dict(st.session_state.get("sp_user") or {})
    for key in (
        "quota_date",
        "daily_free_limit",
        "daily_free_reserved",
        "daily_free_used",
        "daily_free_remaining",
        "bonus_papers",
    ):
        if key in quota:
            user[key] = quota[key]
    if user:
        st.session_state["sp_user"] = user
        st.session_state["sp_account_revision"] = int(
            st.session_state.get("sp_account_revision") or 0
        ) + 1


def _legacy_product_ui_projection() -> dict[str, Any]:
    return {
        "schema_version": "a5-product-ui-v1",
        "mode": "legacy",
        "product_tier_selection_required": False,
        "products": {},
    }


def _normalize_product_ui_projection(value: Any) -> dict[str, Any]:
    """Accept only the Worker-owned A5 UI projection; never reinterpret legacy quota fields."""
    if not isinstance(value, dict):
        return _legacy_product_ui_projection()
    mode = str(value.get("mode") or "legacy").strip().lower()
    if mode == "legacy":
        return _legacy_product_ui_projection()
    if mode != "a5_product":
        raise RuntimeError("product_ui_projection_invalid_mode")

    raw_products = value.get("products")
    raw_entitlement = value.get("entitlement")
    if not isinstance(raw_products, dict) or not isinstance(raw_entitlement, dict):
        raise RuntimeError("product_ui_projection_incomplete")

    products: dict[str, dict[str, Any]] = {}
    for tier in ("STANDARD", "FLAGSHIP"):
        raw = raw_products.get(tier)
        if not isinstance(raw, dict):
            raise RuntimeError("product_ui_projection_missing_tier")
        try:
            minimum = int(raw.get("min_body_words"))
            maximum = int(raw.get("max_body_words"))
        except (TypeError, ValueError) as exc:
            raise RuntimeError("product_ui_projection_invalid_word_policy") from exc
        if minimum <= 0 or maximum < minimum:
            raise RuntimeError("product_ui_projection_invalid_word_policy")
        products[tier] = {
            "product_tier": tier,
            "product_policy_version": str(raw.get("product_policy_version") or ""),
            "word_count_policy_version": str(raw.get("word_count_policy_version") or ""),
            "body_metric": str(raw.get("body_metric") or ""),
            "min_body_words": minimum,
            "max_body_words": maximum,
            "counting_contract_version": str(raw.get("counting_contract_version") or ""),
        }
        if not all((
            products[tier]["product_policy_version"],
            products[tier]["word_count_policy_version"],
            products[tier]["body_metric"],
            products[tier]["counting_contract_version"],
        )):
            raise RuntimeError("product_ui_projection_incomplete")

    try:
        standard_remaining = int(raw_entitlement.get("standard_remaining"))
        flagship_remaining = int(raw_entitlement.get("flagship_remaining"))
    except (TypeError, ValueError) as exc:
        raise RuntimeError("product_ui_projection_invalid_entitlement") from exc
    if standard_remaining < 0 or flagship_remaining < 0:
        raise RuntimeError("product_ui_projection_invalid_entitlement")

    return {
        "schema_version": str(value.get("schema_version") or "a5-product-ui-v1"),
        "mode": "a5_product",
        "product_tier_selection_required": True,
        "products": products,
        "entitlement": {
            "standard_remaining": standard_remaining,
            "flagship_remaining": flagship_remaining,
        },
    }


def get_product_ui_projection() -> dict[str, Any]:
    value = st.session_state.get("sp_product_ui_projection")
    return _normalize_product_ui_projection(value)


def get_account_snapshot(*, block: bool = False) -> dict[str, Any] | None:
    bucket = int(time.time() // 30)
    revision = int(st.session_state.get("sp_account_revision") or 0)
    command_id = f"account:{revision}:{bucket}"
    result = _edge_call("account", command_id, {}, key_prefix="sp_account")
    if result is None:
        if block:
            st.stop()
        return None
    if not result.get("ok"):
        if block:
            _raise_edge_error(result, "account_query_failed")
        return None
    user = result.get("user")
    if isinstance(user, dict):
        st.session_state["sp_user"] = dict(user)
    projection = _normalize_product_ui_projection(result.get("product_ui"))
    st.session_state["sp_product_ui_projection"] = projection
    if projection["mode"] != "a5_product":
        st.session_state.pop("sp_product_tier", None)
    # R3.1 single-writer rule: account refresh owns identity/quota only.
    # Task UI state is initialized by recover_generation_task() and thereafter
    # updated only by generation_task_status()/explicit task actions.
    return result


def start_generation_task(
    *,
    idempotency_key: str,
    mode: str,
    interface_language: str,
    target_words: int,
    product_tier: str | None = None,
) -> dict[str, Any]:
    normalized = str(idempotency_key or "").strip()
    if not normalized:
        raise RuntimeError("idempotency_key_required")
    projection = get_product_ui_projection()
    normalized_tier = str(product_tier or "").strip().upper()
    product_mode = projection.get("mode") == "a5_product"
    if product_mode:
        if normalized_tier not in projection.get("products", {}):
            raise RuntimeError("product_tier_required")
    elif normalized_tier:
        raise RuntimeError("product_tier_not_available_in_legacy_mode")
    payload = {
        "idempotency_key": normalized,
        "mode": str(mode or "").strip().lower(),
        "interface_language": str(interface_language or "zh").strip(),
        "target_words": int(target_words),
    }
    if product_mode:
        # Product tier is user intent only. Worker/D1 resolves entitlement and
        # freezes the trusted TaskPolicySnapshot; browser state is never authority.
        payload["product_tier"] = normalized_tier
    result = _edge_call(
        "task_start",
        normalized,
        payload,
        key_prefix="sp_task_start",
    )
    if result is None:
        st.stop()
    if not result.get("ok"):
        _raise_edge_error(result, "task_start_failed")
    _merge_quota_into_session(result)
    started_task = result.get("task")
    if isinstance(started_task, dict) and started_task.get("task_id"):
        st.session_state["sp_tracked_task_id"] = str(started_task["task_id"])
    return dict(result)


def recover_generation_task() -> dict[str, Any]:
    """Initialize current task state once from the authoritative Worker.

    Account polling is intentionally forbidden from writing task state. This
    recovery call is the only fresh-session initializer; later changes come
    from generation_task_status() or explicit start/cancel operations.
    """
    command_id = f"recover:{int(time.time() // 10)}"
    result = _edge_call("task_recover", command_id, {}, key_prefix="sp_task_recover")
    if result is None:
        st.stop()
    if not result.get("ok"):
        _raise_edge_error(result, "task_recover_failed")
    _merge_quota_into_session(result)
    task = result.get("task")
    if isinstance(task, dict) and task.get("task_id"):
        st.session_state["sp_generation_task"] = dict(task)
        st.session_state["sp_tracked_task_id"] = str(task["task_id"])
        st.session_state["sp_last_ui_task_status"] = str(task.get("status") or "")
    else:
        st.session_state["sp_generation_task"] = None
        st.session_state["sp_last_ui_task_status"] = ""
        st.session_state.pop("sp_tracked_task_id", None)
    return dict(result)


def generation_task_status(task_id: str) -> dict[str, Any]:
    normalized = str(task_id or "").strip()
    command_id = f"status:{normalized}:{int(time.time() // 5)}"
    result = _edge_call(
        "task_status", command_id, {"task_id": normalized}, key_prefix="sp_task_status"
    )
    if result is None:
        st.stop()
    if not result.get("ok"):
        _raise_edge_error(result, "task_status_failed")
    _merge_quota_into_session(result)
    task = result.get("task")
    if isinstance(task, dict):
        st.session_state["sp_generation_task"] = dict(task)
        if task.get("task_id"):
            st.session_state["sp_tracked_task_id"] = str(task["task_id"])
    return dict(result)


def cancel_generation_task(task_id: str | None = None) -> dict[str, Any]:
    """Cancel synchronously so quota release is confirmed before local cleanup."""
    normalized = str(task_id or "").strip()
    command_id = f"cancel:{normalized or 'latest'}"
    result = _edge_call(
        "task_cancel",
        command_id,
        {"task_id": normalized} if normalized else {},
        key_prefix="sp_task_cancel",
    )
    if result is None:
        st.stop()
    if not result.get("ok"):
        _raise_edge_error(result, "task_cancel_failed")
    reservation = result.get("reservation")
    if result.get("cancelled") and reservation and str(reservation.get("status") or "").lower() != "released":
        raise RuntimeError("task_cancel_quota_not_released")
    _merge_quota_into_session(result)
    return dict(result)


def logout_alipay_session(on_logout: Callable[[str, str], None] | None = None) -> None:
    user_id = str((st.session_state.get("sp_user") or {}).get("id") or "")
    task = st.session_state.get("sp_generation_task") or {}
    task_id = str(task.get("task_id") or "")
    if on_logout is not None:
        try:
            on_logout(user_id, task_id)
        except Exception:
            pass
    _clear_session()
    _clear_generation_session_state()


def create_payment_order(idempotency_key: str) -> dict[str, Any]:
    normalized = str(idempotency_key or "").strip()
    result = _edge_call(
        "pay_create",
        normalized,
        {"idempotency_key": normalized},
        key_prefix="sp_pay_create",
    )
    if result is None:
        st.stop()
    if not result.get("ok"):
        _raise_edge_error(result, "payment_create_failed")
    return dict(result)


def payment_order_status(merchant_order_no: str) -> dict[str, Any]:
    order_no = str(merchant_order_no or "").strip()
    command_id = f"pay-status:{order_no}:{int(time.time() // 3)}"
    result = _edge_call(
        "pay_status",
        command_id,
        {"merchant_order_no": order_no},
        key_prefix="sp_pay_status",
    )
    if result is None:
        st.stop()
    if not result.get("ok"):
        _raise_edge_error(result, "payment_status_failed")
    return dict(result)


def _qr_png_bytes(value: str) -> bytes:
    image = qrcode.make(str(value))
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return buffer.getvalue()


def _render_payment_purchase_body(language: str) -> None:
    st.warning(tr(language, "payment_notice"))

    if st.button(tr(language, "payment_create"), key="sp_create_payment_order"):
        st.session_state["sp_payment_request_id"] = (
            f"pay_{int(time.time() * 1000)}_{secrets.token_urlsafe(12)}"
        )

    request_id = str(st.session_state.get("sp_payment_request_id") or "").strip()
    order = st.session_state.get("sp_payment_order")
    if request_id and not order:
        st.info(tr(language, "payment_creating"))
        order_result = create_payment_order(request_id)
        order = dict(order_result.get("order") or {})
        if not order.get("merchant_order_no") or not order.get("qr_code"):
            raise RuntimeError("payment_order_missing_qr_code")
        st.session_state["sp_payment_order"] = order

    order = st.session_state.get("sp_payment_order") or {}
    qr_code = str(order.get("qr_code") or "").strip()
    if qr_code:
        st.image(_qr_png_bytes(qr_code), width=280)
        st.caption(tr(language, "payment_scan"))
        if st.button(tr(language, "payment_check"), key="sp_check_payment_status"):
            status_result = payment_order_status(order.get("merchant_order_no", ""))
            updated = dict(status_result.get("order") or {})
            st.session_state["sp_payment_order"] = updated
            if updated.get("status") == "paid" and updated.get("credit_applied"):
                user = status_result.get("user")
                if isinstance(user, dict):
                    st.session_state["sp_user"] = dict(user)
                    st.session_state["sp_account_revision"] = int(
                        st.session_state.get("sp_account_revision") or 0
                    ) + 1
                st.session_state["sp_payment_required"] = False
                st.session_state["sp_payment_order"] = None
                st.session_state["sp_payment_request_id"] = ""
                st.success(tr(language, "payment_paid"))
            elif updated.get("status") in {"closed", "failed"}:
                st.session_state["sp_payment_order"] = None
                st.session_state["sp_payment_request_id"] = ""
                st.warning(tr(language, "payment_expired"))


def render_payment_purchase(language: str) -> None:
    """Show payment in a modal so the approved normal-page layout is unchanged."""
    dialog_factory = getattr(st, "dialog", None)
    if callable(dialog_factory):
        @dialog_factory(tr(language, "payment_dialog_title"))
        def _payment_dialog():
            _render_payment_purchase_body(language)
        _payment_dialog()
        return
    # Compatibility fallback for older Streamlit builds.
    _render_payment_purchase_body(language)


def render_logged_in_account(
    language: str,
    user: dict[str, Any] | None = None,
    on_logout: Callable[[str, str], None] | None = None,
) -> None:
    """Render the approved responsive account summary with live quota values."""
    live = get_account_snapshot(block=False)
    if live and isinstance(live.get("user"), dict):
        user = dict(live["user"])
    else:
        user = user or st.session_state.get("sp_user", {})
    user_id = user.get("id", "—")
    projection = get_product_ui_projection()
    product_mode = projection.get("mode") == "a5_product"
    if product_mode:
        entitlement = projection["entitlement"]
        standard_remaining = entitlement["standard_remaining"]
        flagship_remaining = entitlement["flagship_remaining"]
    else:
        free_remaining = user.get("daily_free_remaining", user.get("daily_free_limit", 3))
        bonus_papers = user.get("bonus_papers", 0)

    st.markdown(
        """
        <style>
        section[data-testid="stSidebar"] {
            display: none !important;
        }

        [data-testid="collapsedControl"] {
            display: none !important;
        }

        .sp-account-card {
            display: flex;
            align-items: center;
            justify-content: center;
            gap: clamp(0.3rem, 0.55vw, 0.58rem);
            width: 100%;
            min-width: 0;
            min-height: 2.85rem;
            padding: 0.48rem clamp(0.55rem, 0.85vw, 0.9rem);
            border: 1px solid rgba(148, 163, 184, 0.30);
            border-radius: 13px;
            background: rgba(30, 41, 59, 0.58);
            color: var(--text-color, #f3f4f6);
            font-size: clamp(0.76rem, 0.84vw, 0.92rem);
            font-weight: 500;
            line-height: 1.25;
            white-space: nowrap;
            overflow: visible;
            box-sizing: border-box;
        }

        .sp-account-card span {
            display: inline-block;
            white-space: nowrap;
        }

        .sp-account-separator {
            flex: 0 0 auto;
            opacity: 0.48;
        }

        div[data-testid="stMarkdownContainer"]:has(.sp-logout-link) {
            width: 100%;
        }

        div[data-testid="stMarkdownContainer"]:has(.sp-logout-link) p {
            width: 100%;
            margin: 0;
        }

        .sp-logout-link {
            display: flex;
            align-items: center;
            justify-content: center;
            width: 100%;
            min-width: 6.8rem;
            min-height: 2.85rem;
            padding: 0.45rem 0.72rem;
            border: 1px solid rgba(148, 163, 184, 0.30);
            border-radius: 0.5rem;
            background: transparent;
            color: var(--text-color, #f3f4f6) !important;
            text-decoration: none !important;
            white-space: nowrap !important;
            overflow-wrap: normal !important;
            word-break: keep-all !important;
            line-height: 1.15;
            font-size: clamp(0.78rem, 0.82vw, 0.92rem);
            font-weight: 400;
            box-sizing: border-box;
            cursor: pointer;
            transition: border-color 120ms ease, background-color 120ms ease;
        }

        .sp-logout-link:hover {
            border-color: rgba(255, 75, 75, 0.78);
            background: rgba(255, 255, 255, 0.035);
            color: var(--text-color, #f3f4f6) !important;
        }

        .sp-logout-link:focus-visible {
            outline: 2px solid rgba(255, 75, 75, 0.85);
            outline-offset: 2px;
        }

        @media (max-width: 1150px) {
            .sp-account-card {
                font-size: 0.75rem;
                gap: 0.3rem;
                padding-left: 0.48rem;
                padding-right: 0.48rem;
            }

            .sp-logout-link {
                min-width: 5.8rem;
                font-size: 0.78rem;
                padding-left: 0.46rem;
                padding-right: 0.46rem;
            }
        }

        @media (max-width: 760px) {
            .sp-account-card {
                justify-content: flex-start;
                overflow-x: auto;
                scrollbar-width: none;
            }

            .sp-account-card::-webkit-scrollbar {
                display: none;
            }
        }
        </style>
        """,
        unsafe_allow_html=True,
    )

    info_col, logout_col = st.columns([5.9, 1.65])
    with info_col:
        if product_mode:
            balance_markup = (
                f"<span>{html.escape(tr(language, 'account_chip_standard', value=standard_remaining))}</span>"
                f'<span class="sp-account-separator">•</span>'
                f"<span>{html.escape(tr(language, 'account_chip_flagship', value=flagship_remaining))}</span>"
            )
        else:
            balance_markup = (
                f"<span>{html.escape(tr(language, 'account_chip_free', value=free_remaining))}</span>"
                f'<span class="sp-account-separator">•</span>'
                f"<span>{html.escape(tr(language, 'account_chip_bonus', value=bonus_papers))}</span>"
            )
        st.markdown(
            f"""
            <div class="sp-account-card">
              <span>{html.escape(tr(language, 'account_chip_user', value=user_id))}</span>
              <span class="sp-account-separator">•</span>
              {balance_markup}
            </div>
            """,
            unsafe_allow_html=True,
        )
    with logout_col:
        logout_label = html.escape(tr(language, "logout_short"))
        logout_help = html.escape(tr(language, "logout"), quote=True)
        st.markdown(
            f'<a class="sp-logout-link" href="/edge/logout" target="_self" '
            f'rel="nofollow" role="button" aria-label="{logout_help}" '
            f'title="{logout_help}">{logout_label}</a>',
            unsafe_allow_html=True,
        )