deerflow-code/offline-backend-20260512/backend/app/gateway/weknora_embed.py

"""Shared helpers for WeKnora iframe proxy sessions."""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
import secrets
import time
from dataclasses import dataclass
from typing import Any

from fastapi import HTTPException

from deerflow.config.runtime_paths import runtime_home

WEKNORA_EMBED_COOKIE = "deerflow_weknora_embed"
WEKNORA_EMBED_TTL_SECONDS = 15 * 60
WEKNORA_EMBED_API_PREFIX = "/api/llmwiki/weknora-embed"
WEKNORA_ROOT_ASSET_PATHS: frozenset[str] = frozenset(
    {
        "/favicon.ico",
        "/favicon.svg",
        "/config.js",
        "/logo.svg",
        "/manifest.webmanifest",
    }
)
WEKNORA_STATIC_PATH_PREFIXES: tuple[str, ...] = (
    "/assets/",
    "/css/",
    "/fonts/",
    "/img/",
    "/images/",
    "/js/",
    "/locales/",
    "/media/",
    "/static/",
    "/tdesign-icons/",
)


def _load_weknora_proxy_secret() -> str:
    if secret := os.getenv("DEERFLOW_LLMWIKI_PROXY_SECRET"):
        return secret
    secret_path = runtime_home() / "weknora_embed_secret"
    try:
        if secret_path.is_file():
            secret = secret_path.read_text(encoding="utf-8").strip()
            if secret:
                return secret
        secret_path.parent.mkdir(parents=True, exist_ok=True)
        secret = secrets.token_urlsafe(48)
        try:
            fd = os.open(str(secret_path), os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        except FileExistsError:
            return secret_path.read_text(encoding="utf-8").strip()
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            handle.write(secret)
        return secret
    except OSError:
        return secrets.token_urlsafe(48)


WEKNORA_PROXY_SECRET = _load_weknora_proxy_secret()


@dataclass(frozen=True)
class WeKnoraEmbedContext:
    mapping_id: str
    weknora_id: str
    user_id: str
    is_admin: bool
    can_write: bool
    exp: int
    is_conversation_deposit: bool = False
    allowed_weknora_ids: tuple[str, ...] = ()


def is_weknora_embed_proxy_path(path: str) -> bool:
    path = path.removeprefix("/deerflow")
    return (
        path == "/platform"
        or path.startswith("/platform/")
        or path.startswith(WEKNORA_STATIC_PATH_PREFIXES)
        or path in WEKNORA_ROOT_ASSET_PATHS
        or path.startswith(f"{WEKNORA_EMBED_API_PREFIX}/api/v1/")
        or (path.startswith("/api/v1/") and not path.startswith("/api/v1/auth/"))
    )


def _b64url_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")


def _b64url_decode(data: str) -> bytes:
    padded = data + ("=" * (-len(data) % 4))
    return base64.urlsafe_b64decode(padded.encode("ascii"))


def sign_weknora_embed_payload(payload: dict[str, Any]) -> str:
    body = _b64url_encode(json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8"))
    sig = hmac.new(WEKNORA_PROXY_SECRET.encode("utf-8"), body.encode("ascii"), hashlib.sha256).digest()
    return f"{body}.{_b64url_encode(sig)}"


def verify_weknora_signed_payload(raw: str | None) -> dict[str, Any]:
    """Verify a short-lived capability signed with the iframe proxy secret."""

    if not raw or "." not in raw:
        raise HTTPException(status_code=401, detail="Signed resource token is missing")
    body, sig = raw.rsplit(".", 1)
    expected = _b64url_encode(hmac.new(WEKNORA_PROXY_SECRET.encode("utf-8"), body.encode("ascii"), hashlib.sha256).digest())
    if not hmac.compare_digest(sig, expected):
        raise HTTPException(status_code=401, detail="Signed resource token is invalid")
    try:
        payload = json.loads(_b64url_decode(body))
    except (ValueError, TypeError, json.JSONDecodeError):
        raise HTTPException(status_code=401, detail="Signed resource token is invalid") from None
    if not isinstance(payload, dict):
        raise HTTPException(status_code=401, detail="Signed resource token is invalid")
    if int(payload.get("exp") or 0) < int(time.time()):
        raise HTTPException(status_code=401, detail="Signed resource token has expired")
    return payload


def verify_weknora_embed_cookie(raw: str | None) -> WeKnoraEmbedContext:
    if not raw or "." not in raw:
        raise HTTPException(status_code=401, detail="WeKnora iframe session is missing")
    body, sig = raw.rsplit(".", 1)
    expected = _b64url_encode(hmac.new(WEKNORA_PROXY_SECRET.encode("utf-8"), body.encode("ascii"), hashlib.sha256).digest())
    if not hmac.compare_digest(sig, expected):
        raise HTTPException(status_code=401, detail="WeKnora iframe session is invalid")
    try:
        payload = json.loads(_b64url_decode(body))
    except (ValueError, TypeError, json.JSONDecodeError):
        raise HTTPException(status_code=401, detail="WeKnora iframe session is invalid") from None
    if int(payload.get("exp") or 0) < int(time.time()):
        raise HTTPException(status_code=401, detail="WeKnora iframe session has expired")
    raw_allowed_ids = payload.get("allowed_weknora_ids")
    allowed_weknora_ids = tuple(dict.fromkeys(str(item).strip() for item in raw_allowed_ids if isinstance(item, (str, int)) and str(item).strip())) if isinstance(raw_allowed_ids, list) else ()
    return WeKnoraEmbedContext(
        mapping_id=str(payload.get("mapping_id") or ""),
        weknora_id=str(payload.get("weknora_id") or ""),
        user_id=str(payload.get("user_id") or ""),
        is_admin=bool(payload.get("is_admin")),
        can_write=bool(payload.get("can_write")),
        exp=int(payload.get("exp") or 0),
        is_conversation_deposit=bool(payload.get("is_conversation_deposit")),
        allowed_weknora_ids=allowed_weknora_ids,
    )


def has_valid_weknora_embed_cookie(raw: str | None) -> bool:
    try:
        verify_weknora_embed_cookie(raw)
    except HTTPException:
        return False
    return True