xiaobai-review/xiaobai-datahub/datahub/siteauth.py

from __future__ import annotations

import hashlib
import hmac
import json
import threading
import time
import urllib.error
import urllib.request
from typing import Any

from datahub.logutil import get_logger

LOGGER = get_logger()

SITE_SESSION_COOKIE = "xiaobai_session"
VERIFY_CACHE_SECONDS = 20.0
BRIDGE_TIMEOUT_SECONDS = 6.0


class SiteBridgeError(RuntimeError):
    """The review site could not answer a bridge call.

    ``status`` carries the review site's HTTP status when it answered with one.
    A 4xx there means the operator's input was rejected (bad API key, invalid
    model id), which must not surface here as a console fault.
    """

    def __init__(self, message: str, status: int = 0) -> None:
        super().__init__(message)
        self.status = status

    @property
    def caller_fault(self) -> bool:
        return 400 <= self.status < 500


class SiteBridge:
    """Service-to-service client for the review site's ``/api/hub-admin/*`` endpoints.

    The shared ``HUB_ADMIN_TOKEN`` is the only credential; the review site
    checks it before any handler runs, so nothing here needs a browser session.
    """

    def __init__(self, base_url: str, token: str, timeout: float = BRIDGE_TIMEOUT_SECONDS) -> None:
        self.base_url = (base_url or "").rstrip("/")
        self.token = token or ""
        self.timeout = timeout

    @property
    def configured(self) -> bool:
        return bool(self.base_url and self.token)

    def call(self, path: str, payload: dict[str, Any] | None = None) -> dict[str, Any]:
        if not self.configured:
            raise SiteBridgeError("主站桥接未配置：请设置 REVIEW_BASE_URL 与 HUB_ADMIN_TOKEN")
        body = json.dumps(payload or {}, ensure_ascii=False).encode("utf-8")
        request = urllib.request.Request(
            f"{self.base_url}{path}",
            data=body,
            method="POST",
            headers={
                "Content-Type": "application/json; charset=utf-8",
                "X-Hub-Admin-Token": self.token,
            },
        )
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                raw = response.read()
        except urllib.error.HTTPError as exc:
            detail = _error_detail(exc.read())
            raise SiteBridgeError(detail or f"主站返回 {exc.code}", exc.code) from exc
        except (urllib.error.URLError, TimeoutError, OSError) as exc:
            raise SiteBridgeError(f"主站不可达：{exc}") from exc
        try:
            parsed = json.loads(raw.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise SiteBridgeError("主站返回的不是合法 JSON") from exc
        if not isinstance(parsed, dict):
            raise SiteBridgeError("主站返回的不是合法 JSON")
        if parsed.get("error"):
            raise SiteBridgeError(str(parsed["error"]))
        return parsed


def _error_detail(raw: bytes) -> str:
    try:
        parsed = json.loads(raw.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError):
        return ""
    if isinstance(parsed, dict) and parsed.get("error"):
        return str(parsed["error"])
    return ""


class SiteAuth:
    """Admin gate for the console: the review site owns accounts, we only verify.

    The console has no accounts of its own. Every request carries the review
    site's ``xiaobai_session`` cookie (same host, different port, so the browser
    sends it), which we hand to the site for verification. Results are cached
    for a few seconds so a page full of panels does not fan out one bridge call
    per request.

    CSRF is stateless: the token is an HMAC of the session token under a server
    secret, so it is unguessable without the secret yet needs no storage and
    stays valid exactly as long as the session does.
    """

    def __init__(self, bridge: SiteBridge, secret: str, cache_seconds: float = VERIFY_CACHE_SECONDS) -> None:
        self.bridge = bridge
        self._secret = (secret or "").encode("utf-8")
        self._cache_seconds = cache_seconds
        self._cache: dict[str, tuple[float, dict[str, Any] | None]] = {}
        self._lock = threading.Lock()

    def verify(self, session_token: str) -> dict[str, Any] | None:
        if not session_token:
            return None
        key = hashlib.sha256(session_token.encode("utf-8")).hexdigest()
        now = time.monotonic()
        with self._lock:
            cached = self._cache.get(key)
            if cached and cached[0] > now:
                return cached[1]
        user = self._verify_remote(session_token)
        with self._lock:
            self._cache[key] = (now + self._cache_seconds, user)
            if len(self._cache) > 256:
                self._prune(now)
        return user

    def _prune(self, now: float) -> None:
        for cached_key in [key for key, (expires, _) in self._cache.items() if expires <= now]:
            self._cache.pop(cached_key, None)

    def _verify_remote(self, session_token: str) -> dict[str, Any] | None:
        payload = self.bridge.call("/api/hub-admin/session", {"session_token": session_token})
        if not payload.get("authenticated"):
            return None
        user = payload.get("user") or {}
        return {
            "id": int(user.get("id") or 0),
            "username": str(user.get("username") or ""),
            "role": str(user.get("role") or "user"),
            "is_admin": bool(user.get("is_admin")),
        }

    def invalidate(self, session_token: str) -> None:
        key = hashlib.sha256(session_token.encode("utf-8")).hexdigest()
        with self._lock:
            self._cache.pop(key, None)

    def csrf_token(self, session_token: str) -> str:
        digest = hashlib.sha256(session_token.encode("utf-8")).digest()
        return hmac.new(self._secret, digest, hashlib.sha256).hexdigest()

    def check_csrf(self, session_token: str, supplied: str) -> bool:
        if not supplied:
            return False
        return hmac.compare_digest(self.csrf_token(session_token), supplied)

    def logout(self, session_token: str) -> None:
        self.invalidate(session_token)
        if not session_token:
            return
        try:
            self.bridge.call("/api/hub-admin/session/logout", {"session_token": session_token})
        except SiteBridgeError:
            LOGGER.warning("site logout bridge call failed")

    def confirm_password(self, user_id: int, password: str) -> bool:
        if not password:
            return False
        payload = self.bridge.call(
            "/api/hub-admin/password/check",
            {"user_id": int(user_id), "password": password},
        )
        return bool(payload.get("verified"))