zcbot/core/external_systems/auth.py

"""外部系统认证策略。

认证只消费管理员保存的可信配置和用户加密保存的字段，不允许模型指定认证地址或请求头。
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any, Protocol
from urllib.parse import urljoin

import httpx


class ExternalAuthError(RuntimeError):
    pass


@dataclass(frozen=True)
class CredentialField:
    name: str
    label: str
    secret: bool = True
    autocomplete: str = "off"

    def as_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "label": self.label,
            "secret": self.secret,
            "autocomplete": self.autocomplete,
        }


class AuthStrategy(Protocol):
    credential_fields: tuple[CredentialField, ...]

    def headers(
        self,
        *,
        client: httpx.Client,
        base_url: str,
        credentials: dict[str, str],
        config: dict[str, Any],
    ) -> dict[str, str]: ...


def _required(credentials: dict[str, str], fields: tuple[CredentialField, ...]) -> None:
    missing = [
        field.label for field in fields if not credentials.get(field.name, "").strip()
    ]
    if missing:
        raise ExternalAuthError("请填写" + "、".join(missing))


def _nested_value(payload: Any, path: str) -> Any:
    current = payload
    for part in path.split("."):
        if not isinstance(current, dict):
            return None
        current = current.get(part)
    return current


def _auth_header(
    config: dict[str, Any], token: str, *, default_name: str, default_template: str
) -> dict[str, str]:
    name = str(config.get("auth_header_name") or default_name).strip()
    template = str(config.get("auth_header_template") or default_template)
    if any(char in name for char in "\r\n:") or any(
        char in template for char in "\r\n"
    ):
        raise ExternalAuthError("认证 Header 配置非法")
    if "{token}" not in template:
        raise ExternalAuthError("认证 Header 模板必须包含 {token}")
    return {name: template.replace("{token}", token)}


class PasswordJwtAuth:
    credential_fields: tuple[CredentialField, ...] = (
        CredentialField("username", "用户名", secret=False, autocomplete="username"),
        CredentialField("password", "密码", autocomplete="current-password"),
    )

    def headers(
        self,
        *,
        client: httpx.Client,
        base_url: str,
        credentials: dict[str, str],
        config: dict[str, Any],
    ) -> dict[str, str]:
        _required(credentials, self.credential_fields)
        login_path = str(config.get("login_path") or "/api/auth/token/").strip()
        if not login_path.startswith("/") or "://" in login_path:
            raise ExternalAuthError("login_path 必须是站内绝对路径")
        username_field = str(config.get("username_field") or "username").strip()
        password_field = str(config.get("password_field") or "password").strip()
        token_field = str(config.get("token_field") or "access").strip()
        try:
            with client.stream(
                "POST",
                urljoin(base_url + "/", login_path.lstrip("/")),
                json={
                    username_field: credentials["username"],
                    password_field: credentials["password"],
                },
            ) as response:
                chunks: list[bytes] = []
                total = 0
                for chunk in response.iter_bytes():
                    total += len(chunk)
                    if total > 65536:
                        raise ExternalAuthError("外部系统登录响应超过安全上限")
                    chunks.append(chunk)
                status_code = response.status_code
                content = b"".join(chunks)
        except httpx.HTTPError as exc:
            raise ExternalAuthError(
                f"外部系统登录连接失败: {type(exc).__name__}"
            ) from exc
        if status_code >= 400:
            raise ExternalAuthError(f"外部系统登录失败(HTTP {status_code})")
        try:
            token = _nested_value(json.loads(content.decode("utf-8-sig")), token_field)
        except (UnicodeDecodeError, ValueError):
            token = None
        if not isinstance(token, str) or not token:
            raise ExternalAuthError(f"外部系统登录响应缺少 {token_field}")
        return _auth_header(
            config,
            token,
            default_name="Authorization",
            default_template="Bearer {token}",
        )


class ApiKeyAuth:
    credential_fields: tuple[CredentialField, ...] = (
        CredentialField("api_key", "API Key"),
    )

    def headers(
        self,
        *,
        client: httpx.Client,
        base_url: str,
        credentials: dict[str, str],
        config: dict[str, Any],
    ) -> dict[str, str]:
        _required(credentials, self.credential_fields)
        return _auth_header(
            config,
            credentials["api_key"],
            default_name="X-API-Key",
            default_template="{token}",
        )


class BearerTokenAuth:
    credential_fields: tuple[CredentialField, ...] = (
        CredentialField("token", "Bearer Token"),
    )

    def headers(
        self,
        *,
        client: httpx.Client,
        base_url: str,
        credentials: dict[str, str],
        config: dict[str, Any],
    ) -> dict[str, str]:
        _required(credentials, self.credential_fields)
        return _auth_header(
            config,
            credentials["token"],
            default_name="Authorization",
            default_template="Bearer {token}",
        )


_AUTH_STRATEGIES: dict[str, AuthStrategy] = {
    "password_jwt": PasswordJwtAuth(),
    "api_key": ApiKeyAuth(),
    "bearer_token": BearerTokenAuth(),
}


def get_auth_strategy(auth_type: str) -> AuthStrategy:
    strategy = _AUTH_STRATEGIES.get((auth_type or "").strip())
    if strategy is None:
        raise ExternalAuthError(f"不支持的认证方式: {auth_type}")
    return strategy


def auth_catalog() -> list[dict[str, Any]]:
    titles = {
        "password_jwt": "用户名密码换取 Token",
        "api_key": "API Key",
        "bearer_token": "Bearer Token",
    }
    return [
        {
            "auth_type": key,
            "title": titles[key],
            "credential_fields": [
                field.as_dict() for field in strategy.credential_fields
            ],
        }
        for key, strategy in _AUTH_STRATEGIES.items()
    ]