Aether/_deprecated_py_src/services/provider_ops/architectures/nekocode.py

"""
NekoCode 架构

针对 NekoCode 中转站的预设配置，使用 Cookie 认证。
"""

from typing import Any

import httpx

from src.core.logger import logger
from src.services.provider_ops.actions import ProviderAction
from src.services.provider_ops.actions.nekocode_balance import NekoCodeBalanceAction
from src.services.provider_ops.architectures.base import (
    ProviderArchitecture,
    ProviderConnector,
    VerifyResult,
)
from src.services.provider_ops.types import ConnectorAuthType, ProviderActionType
from src.services.provider_ops.utils import extract_cookie_value
from src.utils.ssl_utils import get_ssl_context


class NekoCodeConnector(ProviderConnector):
    """
    NekoCode 专用连接器

    特点：
    - 使用 Cookie 认证（session）
    """

    auth_type = ConnectorAuthType.COOKIE
    display_name = "NekoCode Cookie"

    def __init__(self, base_url: str, config: dict[str, Any] | None = None):
        super().__init__(base_url, config)
        self._session_cookie: str | None = None

    async def connect(self, credentials: dict[str, Any]) -> bool:
        """建立连接"""
        session_cookie = credentials.get("session_cookie")
        if not session_cookie:
            self._set_error("Session Cookie 不能为空")
            return False

        # 提取纯 session 值（支持完整 Cookie 字符串或仅 session 值）
        self._session_cookie = extract_cookie_value(session_cookie, "session")

        self._set_connected()
        return True

    async def disconnect(self) -> None:
        """断开连接"""
        self._session_cookie = None
        self._set_disconnected()

    async def is_authenticated(self) -> bool:
        """检查是否已认证"""
        return self._session_cookie is not None

    def _apply_auth(self, request: httpx.Request) -> httpx.Request:
        """为请求应用认证信息"""
        if self._session_cookie:
            request.headers["Cookie"] = f"session={self._session_cookie}"

        return request

    @classmethod
    def get_credentials_schema(cls) -> dict[str, Any]:
        """获取凭据配置 schema"""
        return {
            "type": "object",
            "properties": {
                "base_url": {
                    "type": "string",
                    "title": "站点地址",
                    "description": "API 基础地址",
                    "x-default-value": "https://nekocode.ai",
                },
                "session_cookie": {
                    "type": "string",
                    "title": "Session Cookie",
                    "description": "从浏览器复制的 session Cookie 值",
                    "x-sensitive": True,
                    "x-input-type": "password",
                },
            },
            "required": ["session_cookie"],
            "x-field-groups": [
                {"fields": ["base_url"]},
                {"fields": ["session_cookie"]},
            ],
            "x-auth-type": "cookie",
            "x-default-base-url": "https://nekocode.ai",
            "x-validation": [
                {
                    "type": "required",
                    "fields": ["session_cookie"],
                    "message": "请填写 Session Cookie",
                },
            ],
            "x-quota-divisor": None,
            "x-currency": "USD",
            "x-balance-extra-format": [
                {
                    "label": "天",
                    "type": "daily_quota",
                    "source_limit": "daily_quota_limit",
                    "source_remaining": "daily_remaining_quota",
                    "source_start_date": "effective_start_date",
                },
                {
                    "label": "月",
                    "type": "monthly_expiry",
                    "source_end_date": "effective_end_date",
                },
            ],
        }


class NekoCodeArchitecture(ProviderArchitecture):
    """
    NekoCode 架构预设

    针对 NekoCode 中转站优化的预设配置。

    特点：
    - 使用 Cookie 认证（session）
    - 验证端点: /api/usage/summary
    - 显示余额、每日配额、订阅状态
    """

    architecture_id = "nekocode"
    display_name = "NekoCode"
    description = "NekoCode 中转站预设配置，使用 Cookie 认证"

    supported_connectors: list[type[ProviderConnector]] = [
        NekoCodeConnector,
    ]

    supported_actions: list[type[ProviderAction]] = [NekoCodeBalanceAction]

    default_action_configs: dict[ProviderActionType, dict[str, Any]] = {
        ProviderActionType.QUERY_BALANCE: {
            "endpoint": "/api/usage/summary",
            "method": "GET",
        },
    }

    def get_credentials_schema(self) -> dict[str, Any]:
        """NekoCode 使用 session_cookie 认证"""
        return NekoCodeConnector.get_credentials_schema()

    def get_verify_endpoint(self) -> str:
        """验证端点"""
        return "/api/user/self"

    async def prepare_verify_config(
        self,
        base_url: str,
        config: dict[str, Any],
        credentials: dict[str, Any],
    ) -> dict[str, Any]:
        """
        验证前获取 /api/usage/summary 数据（用于显示天卡信息）

        Args:
            base_url: API 基础地址
            config: 连接器配置
            credentials: 凭据信息

        Returns:
            包含 _usage_summary 的配置（会被合并到验证响应中）
        """
        try:
            # 构建请求头
            headers: dict[str, str] = {
                "User-Agent": (
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                    "AppleWebKit/537.36 (KHTML, like Gecko) "
                    "Chrome/120.0.0.0 Safari/537.36"
                ),
            }

            # 添加 Cookie
            cookie_input = credentials.get("session_cookie")
            if cookie_input:
                session_value = extract_cookie_value(cookie_input, "session")
                headers["Cookie"] = f"session={session_value}"

            # 构建 client 参数
            client_kwargs: dict[str, Any] = {
                "timeout": 10,
                "verify": get_ssl_context(),
            }
            from src.services.proxy_node.resolver import resolve_ops_proxy_config_async

            proxy, tunnel_node_id = await resolve_ops_proxy_config_async(config)
            if tunnel_node_id:
                from src.services.proxy_node.tunnel_transport import create_tunnel_transport

                client_kwargs["transport"] = create_tunnel_transport(tunnel_node_id, timeout=10.0)
            elif proxy:
                client_kwargs["proxy"] = proxy

            async with httpx.AsyncClient(**client_kwargs) as client:
                resp = await client.get(
                    f"{base_url.rstrip('/')}/api/usage/summary",
                    headers=headers,
                )

                if resp.status_code == 200:
                    data = resp.json()
                    if data.get("success"):
                        return {"_usage_summary": data.get("data", {})}

        except Exception as e:
            logger.debug("获取 NekoCode usage summary 失败: {}", e)

        return {}

    def build_verify_headers(
        self,
        config: dict[str, Any],
        credentials: dict[str, Any],
    ) -> dict[str, str]:
        """
        构建 NekoCode 的验证请求 Headers

        使用 Cookie 认证，不使用 Authorization。
        """
        headers: dict[str, str] = {}

        # 添加 session Cookie
        cookie_input = credentials.get("session_cookie")
        if cookie_input:
            # 提取 session 值（支持完整 Cookie 字符串或仅 session 值）
            session_value = extract_cookie_value(cookie_input, "session")
            headers["Cookie"] = f"session={session_value}"

        return headers

    def _auth_fail_message(self, status_code: int) -> str:
        """Cookie 认证的错误消息"""
        if status_code == 401:
            return "Cookie 已失效，请重新配置"
        return "Cookie 已失效或无权限"

    def _build_verify_result(
        self, user_data: dict[str, Any], raw_data: dict[str, Any] | None = None
    ) -> VerifyResult:
        """NekoCode 自定义字段提取（合并 _usage_summary 天卡数据）"""
        # 转换余额字符串为数字
        balance = user_data.get("balance")
        try:
            quota = float(balance) if balance else None
        except (TypeError, ValueError):
            quota = None

        # 从 prepare_verify_config 获取的 _usage_summary 数据（天卡信息）
        extra: dict[str, Any] = {}
        usage_summary = (raw_data or {}).get("_usage_summary", {})
        subscription = usage_summary.get("subscription", {})

        if subscription:
            daily_limit = subscription.get("daily_quota_limit")
            daily_remaining = subscription.get("daily_remaining_quota")
            try:
                extra["daily_quota_limit"] = float(daily_limit) if daily_limit else None
            except (TypeError, ValueError):
                pass
            try:
                extra["daily_remaining_quota"] = float(daily_remaining) if daily_remaining else None
            except (TypeError, ValueError):
                pass

            extra["plan_name"] = subscription.get("plan_name")
            extra["subscription_status"] = subscription.get("status")
            extra["effective_start_date"] = subscription.get("effective_start_date")
            extra["effective_end_date"] = subscription.get("effective_end_date")

        return VerifyResult(
            success=True,
            username=user_data.get("username"),
            display_name=user_data.get("display_name") or user_data.get("username"),
            email=user_data.get("email"),
            quota=quota,
            extra=extra if extra else None,
        )