Files

103 lines
3.3 KiB
Python

"""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
WEKNORA_EMBED_COOKIE = "deerflow_weknora_embed"
WEKNORA_EMBED_TTL_SECONDS = 15 * 60
WEKNORA_EMBED_API_PREFIX = "/api/llmwiki/weknora-embed"
WEKNORA_PROXY_SECRET = os.getenv("DEERFLOW_LLMWIKI_PROXY_SECRET") or secrets.token_urlsafe(48)
WEKNORA_ROOT_ASSET_PATHS: frozenset[str] = frozenset(
{
"/favicon.ico",
"/favicon.svg",
"/config.js",
"/logo.svg",
"/manifest.webmanifest",
}
)
@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
def is_weknora_embed_proxy_path(path: str) -> bool:
return (
path == "/platform"
or path.startswith("/platform/")
or path.startswith("/assets/")
or path.startswith("/locales/")
or path.startswith("/tdesign-icons/")
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_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")
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")),
)
def has_valid_weknora_embed_cookie(raw: str | None) -> bool:
try:
verify_weknora_embed_cookie(raw)
except HTTPException:
return False
return True