""" YesCode 架构 针对 YesCode 中转站的预设配置。 """ from typing import Any import httpx from src.services.provider_ops.actions import ProviderAction from src.services.provider_ops.actions.yescode_balance import ( YesCodeBalanceAction, fetch_yescode_combined_data, parse_yescode_balance_extra, ) from src.services.provider_ops.architectures.base import ( ProviderArchitecture, ProviderConnector, VerifyResult, ) from src.services.provider_ops.types import ConnectorAuthType, ProviderActionType from src.utils.ssl_utils import get_ssl_context def _extract_cookies(cookie_string: str) -> dict[str, str]: """ 从完整的 Cookie 字符串中提取 yescode_auth 和 yescode_csrf Args: cookie_string: Cookie 字符串 Returns: 包含 yescode_auth 和 yescode_csrf 的字典 """ result: dict[str, str] = {} for part in cookie_string.split(";"): part = part.strip() if "=" in part: key, value = part.split("=", 1) key = key.strip() if key in ("yescode_auth", "yescode_csrf"): result[key] = value.strip() return result def _build_cookie_header(cookie_string: str) -> str: """ 从输入的 Cookie 字符串构建请求用的 Cookie header 支持两种输入格式: 1. 完整 Cookie: "yescode_auth=xxx; yescode_csrf=yyy" 2. 仅 auth 值: "eyJhbGciOiJI..." Args: cookie_string: Cookie 字符串或 auth 值 Returns: Cookie header 值 """ # 如果包含 "yescode_auth=",说明是完整 Cookie 字符串 if "yescode_auth=" in cookie_string: cookies = _extract_cookies(cookie_string) parts = [] if "yescode_auth" in cookies: parts.append(f"yescode_auth={cookies['yescode_auth']}") if "yescode_csrf" in cookies: parts.append(f"yescode_csrf={cookies['yescode_csrf']}") return "; ".join(parts) # 否则认为直接是 auth 值 return f"yescode_auth={cookie_string.strip()}" class YesCodeConnector(ProviderConnector): """ YesCode 专用连接器 特点: - 使用 Cookie 认证(yescode_auth JWT + yescode_csrf) """ auth_type = ConnectorAuthType.COOKIE display_name = "YesCode Cookie" def __init__(self, base_url: str, config: dict[str, Any] | None = None): super().__init__(base_url, config) self._auth_cookie: str | None = None async def connect(self, credentials: dict[str, Any]) -> bool: """建立连接""" auth_cookie = credentials.get("auth_cookie") if not auth_cookie: self._set_error("Auth Cookie 不能为空") return False # 构建 Cookie header self._auth_cookie = _build_cookie_header(auth_cookie) self._set_connected() return True async def disconnect(self) -> None: """断开连接""" self._auth_cookie = None self._set_disconnected() async def is_authenticated(self) -> bool: """检查是否已认证""" return self._auth_cookie is not None def _apply_auth(self, request: httpx.Request) -> httpx.Request: """为请求应用认证信息""" if self._auth_cookie: request.headers["Cookie"] = self._auth_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://co.yes.vg", }, "auth_cookie": { "type": "string", "title": "Auth Cookie", "description": "从浏览器复制的 Cookie(包含 yescode_auth 和 yescode_csrf)", "x-sensitive": True, "x-input-type": "password", }, }, "required": ["auth_cookie"], "x-field-groups": [ {"fields": ["base_url"]}, {"fields": ["auth_cookie"]}, ], "x-auth-type": "cookie", "x-default-base-url": "https://co.yes.vg", "x-validation": [ { "type": "required", "fields": ["auth_cookie"], "message": "请填写 Auth Cookie", }, ], "x-quota-divisor": None, "x-currency": "USD", "x-balance-extra-format": [ { "label": "天", "type": "weekly_spent", "source_limit": "daily_limit", "source_spent": "daily_spent", "source_resets_at": "daily_resets_at", }, { "label": "周", "type": "weekly_spent", "source_limit": "weekly_limit", "source_spent": "weekly_spent", "source_resets_at": "weekly_resets_at", }, ], } class YesCodeArchitecture(ProviderArchitecture): """ YesCode 架构预设 针对 YesCode 中转站优化的预设配置。 特点: - 使用 Cookie 认证(yescode_auth JWT + yescode_csrf) - 验证端点: /api/v1/user/balance - 余额单位直接是美元 - 支持每周限额查询 """ architecture_id = "yescode" display_name = "YesCode" description = "YesCode 中转站预设配置,使用 Cookie 认证" supported_connectors: list[type[ProviderConnector]] = [ YesCodeConnector, ] supported_actions: list[type[ProviderAction]] = [ YesCodeBalanceAction, ] default_action_configs: dict[ProviderActionType, dict[str, Any]] = { ProviderActionType.QUERY_BALANCE: { "endpoint": "/api/v1/user/balance", "method": "GET", "response_mapping": { "total_available": "total_balance", }, }, } def get_credentials_schema(self) -> dict[str, Any]: """YesCode 使用 auth_cookie 认证""" return YesCodeConnector.get_credentials_schema() def get_verify_endpoint(self) -> str: """验证端点 - 使用 profile 接口获取完整信息""" return "/api/v1/auth/profile" async def prepare_verify_config( self, base_url: str, config: dict[str, Any], credentials: dict[str, Any], ) -> dict[str, Any]: """ 预获取合并数据(balance + profile) 验证时并发调用两个接口获取完整数据。 """ extra_config: dict[str, Any] = {} cookie_input = credentials.get("auth_cookie") if not cookie_input: return extra_config cookie_header = _build_cookie_header(cookie_input) # 获取代理配置(支持 proxy_node_id、tunnel 和旧的 proxy URL) from src.services.proxy_node.resolver import resolve_ops_proxy_config_async proxy, tunnel_node_id = await resolve_ops_proxy_config_async(config) try: # 构建 client 参数 client_kwargs: dict[str, Any] = { "headers": {"Cookie": cookie_header}, "timeout": 10.0, "verify": get_ssl_context(), } 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 # 创建临时 client 获取合并数据 async with httpx.AsyncClient(**client_kwargs) as client: combined_data = await fetch_yescode_combined_data(client, base_url) extra_config["_combined_data"] = combined_data except Exception: # 如果调用失败,不影响验证流程(会回退到单独调用 profile) pass return extra_config def build_verify_headers( self, config: dict[str, Any], credentials: dict[str, Any], ) -> dict[str, str]: """ 构建 YesCode 的验证请求 Headers 使用 Cookie 认证,不使用 Authorization。 """ headers: dict[str, str] = {} # 添加 Cookie cookie_input = credentials.get("auth_cookie") if cookie_input: headers["Cookie"] = _build_cookie_header(cookie_input) return headers def parse_verify_response( self, status_code: int, data: dict[str, Any], ) -> VerifyResult: """解析 YesCode 验证响应(使用预获取的合并数据)""" if status_code == 401: return VerifyResult(success=False, message="Cookie 已失效,请重新配置") if status_code == 403: return VerifyResult(success=False, message="Cookie 已失效或无权限") if status_code != 200: return VerifyResult(success=False, message=f"验证失败:HTTP {status_code}") # 优先使用预获取的合并数据(包含 balance + profile) combined_data = data.get("_combined_data") if combined_data: # 检查是否有有效数据 if "_profile_data" not in combined_data and "_balance_data" not in combined_data: return VerifyResult(success=False, message="Cookie 已失效,请重新配置") # 使用公共函数解析余额 extra = parse_yescode_balance_extra(combined_data) total_available = extra.pop("_total_available", 0) extra.pop("_subscription_available", None) return VerifyResult( success=True, username=combined_data.get("username"), display_name=combined_data.get("username"), email=combined_data.get("email"), quota=total_available, extra=extra if extra else None, ) # 回退:仅使用 profile 数据(旧逻辑,当 prepare_verify_config 失败时) if "username" not in data: return VerifyResult(success=False, message="响应格式无效") # 构造兼容格式供公共函数使用 compat_data = { "pay_as_you_go_balance": data.get("pay_as_you_go_balance", 0), "subscription_balance": data.get("subscription_balance", 0), "weekly_spent_balance": data.get("current_week_spend", 0), "subscription_plan": data.get("subscription_plan"), "last_week_reset": data.get("last_week_reset"), "last_daily_balance_add": data.get("last_daily_balance_add"), } extra = parse_yescode_balance_extra(compat_data) total_available = extra.pop("_total_available", 0) extra.pop("_subscription_available", None) return VerifyResult( success=True, username=data.get("username"), display_name=data.get("username"), email=data.get("email"), quota=total_available, extra=extra if extra else None, )