173 lines
7.4 KiB
Python
173 lines
7.4 KiB
Python
"""Provider 候选凭据测试与统一、保守的错误分类。"""
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import hmac
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
from collections.abc import Callable
|
||
from dataclasses import dataclass
|
||
from decimal import Decimal, InvalidOperation
|
||
from typing import Any
|
||
|
||
import httpx
|
||
|
||
from .registry import ProviderDefinition, get_provider
|
||
|
||
_EXHAUSTED_RE = re.compile(
|
||
r"insufficient[ _-]?(?:balance|quota|credit)|balance[ _-]?not[ _-]?enough|余额不足|额度不足|quota exhausted",
|
||
re.IGNORECASE,
|
||
)
|
||
_AUTH_RE = re.compile(
|
||
r"invalid.*(?:key|token)|authentication|unauthori[sz]ed|鉴权|认证失败",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class TestResult:
|
||
status: str
|
||
detail: str
|
||
balance_amount: Decimal | None = None
|
||
balance_currency: str | None = None
|
||
|
||
@property
|
||
def accepted(self) -> bool:
|
||
return self.status in {"normal", "low_balance"}
|
||
|
||
|
||
def classify_response(status_code: int, text: str = "") -> TestResult:
|
||
safe = re.sub(r"\s+", " ", text or "").strip()[:300]
|
||
if status_code in {401, 403} or _AUTH_RE.search(safe):
|
||
return TestResult("auth_error", f"认证失败（HTTP {status_code}）")
|
||
if status_code == 402 or _EXHAUSTED_RE.search(safe):
|
||
return TestResult("exhausted", f"余额或额度已耗尽（HTTP {status_code}）")
|
||
if status_code == 429:
|
||
return TestResult("unreachable", "服务限流（HTTP 429），未判定为余额耗尽")
|
||
if status_code >= 400:
|
||
return TestResult("unreachable", f"服务返回 HTTP {status_code}")
|
||
return TestResult("normal", "认证与连通性正常")
|
||
|
||
|
||
def _deepseek(response: httpx.Response, provider: ProviderDefinition) -> TestResult:
|
||
base = classify_response(response.status_code, response.text)
|
||
if base.status != "normal":
|
||
return base
|
||
try:
|
||
body = response.json()
|
||
infos = body.get("balance_infos") or []
|
||
cny = next(x for x in infos if str(x.get("currency", "")).upper() == "CNY")
|
||
amount = Decimal(str(cny.get("total_balance")))
|
||
except (ValueError, KeyError, StopIteration, InvalidOperation, TypeError):
|
||
return TestResult("unreachable", "余额响应格式异常")
|
||
if body.get("is_available") is False or amount <= 0:
|
||
return TestResult("exhausted", "余额已耗尽", amount, "CNY")
|
||
threshold = Decimal(str(provider.low_balance_threshold or 0))
|
||
status = "low_balance" if amount < threshold else "normal"
|
||
detail = f"CNY 可用余额 ¥{amount:.2f}"
|
||
return TestResult(status, detail, amount, "CNY")
|
||
|
||
|
||
def _xfyun_sign(credentials: dict[str, str]) -> dict[str, str]:
|
||
ts = str(int(time.time()))
|
||
md5hex = hashlib.md5(
|
||
(credentials["appid"] + ts).encode(), usedforsecurity=False
|
||
).hexdigest()
|
||
signa = base64.b64encode(
|
||
hmac.new(credentials["secret_key"].encode(), md5hex.encode(), hashlib.sha1).digest()
|
||
).decode()
|
||
return {"appId": credentials["appid"], "ts": ts, "signa": signa,
|
||
"orderId": "zcbot-credential-check", "resultType": "transfer"}
|
||
|
||
|
||
def test_provider(
|
||
provider_id: str,
|
||
credentials: dict[str, str],
|
||
*,
|
||
request: Callable[..., httpx.Response] = httpx.request,
|
||
) -> TestResult:
|
||
provider = get_provider(provider_id)
|
||
expected = {field.name for field in provider.fields}
|
||
if set(credentials) != expected or any(not str(v).strip() for v in credentials.values()):
|
||
return TestResult("auth_error", "凭据字段不完整")
|
||
headers: dict[str, str] = {}
|
||
params: dict[str, str] = {}
|
||
json_body: dict[str, Any] | None = None
|
||
method = "GET"
|
||
test_url = provider.test_url
|
||
if provider.test_url_env:
|
||
base = (os.getenv(provider.test_url_env) or "").strip().rstrip("/")
|
||
if base and provider_id == "document_search":
|
||
test_url = f"{base}/document_search/list_knowledge_bases"
|
||
elif base and provider_id == "paper_server":
|
||
test_url = f"{base}/api/resm/paper/"
|
||
if provider.test_kind in {"deepseek_balance", "bearer_get", "bocha_search"}:
|
||
headers["Authorization"] = f"Bearer {credentials['api_key']}"
|
||
if provider.test_kind == "query_get":
|
||
params["api_key"] = credentials["api_key"]
|
||
params["page_size"] = "1"
|
||
elif provider.test_kind == "mp_get":
|
||
headers["X-API-KEY"] = credentials["api_key"]
|
||
elif provider.test_kind == "bocha_search":
|
||
method = "POST"
|
||
json_body = {"query": "水泥", "count": 1, "freshness": "noLimit"}
|
||
elif provider.test_kind == "xfyun_lfasr":
|
||
method = "POST"
|
||
params.update(_xfyun_sign(credentials))
|
||
elif provider.test_kind == "xfyun_iat":
|
||
# IAT 鉴权只存在于 WebSocket upgrade；复用官方签名函数并允许测试替换
|
||
# requester。200/101 均视为握手成功，真实 handler 不发送音频、不计费。
|
||
from core.asr_xfyun import build_auth_url
|
||
url = build_auth_url(credentials["api_key"], credentials["api_secret"])
|
||
if request is httpx.request:
|
||
try:
|
||
from websockets.sync.client import connect
|
||
with connect(url, open_timeout=8) as websocket:
|
||
websocket.send(json.dumps({
|
||
"common": {"app_id": credentials["appid"]},
|
||
"business": {"language": "zh_cn", "domain": "iat", "accent": "mandarin"},
|
||
"data": {"status": 2, "format": "audio/L16;rate=16000",
|
||
"encoding": "raw", "audio": ""},
|
||
}))
|
||
body = json.loads(websocket.recv(timeout=8))
|
||
if int(body.get("code") or 0) == 0:
|
||
return TestResult("normal", "WebSocket 凭据组鉴权正常")
|
||
return TestResult("auth_error", "WebSocket 凭据组认证失败")
|
||
except Exception as exc: # noqa: BLE001 - WebSocket 库异常族随版本变化
|
||
text = str(exc)
|
||
status = int(getattr(exc, "status_code", 0) or 0)
|
||
return classify_response(status, text) if status else TestResult(
|
||
"unreachable", "WebSocket 连接失败"
|
||
)
|
||
response = request("GET", url, headers={"X-Appid": credentials["appid"]}, timeout=8)
|
||
return classify_response(response.status_code, response.text)
|
||
try:
|
||
response = request(
|
||
method, test_url, headers=headers, params=params,
|
||
json=json_body, timeout=8,
|
||
)
|
||
except httpx.HTTPError:
|
||
return TestResult("unreachable", "网络连接失败")
|
||
if provider.test_kind == "deepseek_balance":
|
||
return _deepseek(response, provider)
|
||
if provider.test_kind == "xfyun_lfasr" and response.status_code < 400:
|
||
try:
|
||
body = response.json()
|
||
code = str(body.get("code") or "")
|
||
description = str(body.get("descInfo") or "")
|
||
if code in {"10105", "10106", "10107"}:
|
||
return TestResult("auth_error", "认证失败")
|
||
# 虚拟订单的“订单不存在”说明签名已通过。
|
||
if code and re.search(
|
||
r"order|订单.*(?:不存在|无效)", description, re.IGNORECASE
|
||
):
|
||
return TestResult("normal", "认证正常（未创建计费订单）")
|
||
if code:
|
||
return TestResult("unreachable", "服务未确认凭据有效")
|
||
except (ValueError, AttributeError):
|
||
pass
|
||
return classify_response(response.status_code, response.text)
|