mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
290 lines
9.7 KiB
Python
290 lines
9.7 KiB
Python
"""
|
||
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,
|
||
)
|