Files
ai-balance-monitor/app/fetcher.py
T

119 lines
4.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""余额获取:模板渲染 → HTTP 请求(带重试)→ JSON 路径提取。"""
from __future__ import annotations
import json
import logging
import re
import time
from dataclasses import dataclass
import requests
logger = logging.getLogger("monitor.fetcher")
PLACEHOLDER = "{{apiKey}}"
_TOKEN_RE = re.compile(r"([^.\[]+)|\[(\d+)\]")
@dataclass
class FetchResult:
"""一次拉取的结果。"""
ok: bool
balance: float | None = None
auth_error: bool = False
status_code: int | None = None
error: str = ""
def render_template(text: str, api_key: str) -> str:
"""把模板中的 {{apiKey}} 替换为账号密钥。"""
return text.replace(PLACEHOLDER, api_key)
def extract_balance(data: object, path: str) -> float:
"""按点路径/数组索引提取余额,如 data.balance、data[0].balance、$.data[0].balance。
取到的值必须是数字或可转数字的字符串,否则抛 ValueError。
"""
p = path.strip()
if p.startswith("$"):
p = p[1:]
cur: object = data
for key, idx in _TOKEN_RE.findall(p):
if key:
if not isinstance(cur, dict) or key not in cur:
raise ValueError(f"路径不存在: {path}(在 {key!r} 处)")
cur = cur[key]
if idx != "":
n = int(idx)
if not isinstance(cur, list) or n >= len(cur):
raise ValueError(f"数组索引越界: {path}(索引 {n}")
cur = cur[n]
if isinstance(cur, bool) or not isinstance(cur, (int, float, str)):
raise ValueError(f"余额不是数字: {path} -> {cur!r}")
try:
return float(cur)
except (TypeError, ValueError):
raise ValueError(f"余额不是数字: {path} -> {cur!r}")
def _build_request(platform: dict, api_key: str) -> tuple[str, dict, str | None]:
url = render_template(platform["url"], api_key)
headers_raw = platform.get("headers") or {}
if isinstance(headers_raw, str):
headers_raw = json.loads(headers_raw) if headers_raw.strip() else {}
headers = json.loads(render_template(json.dumps(headers_raw), api_key))
body = None
if platform["method"] == "POST" and platform.get("body"):
body = render_template(platform["body"], api_key)
return url, headers, body
def fetch_balance(platform: dict, api_key: str, retry_count: int, timeout: int) -> FetchResult:
"""执行一次余额拉取。
策略:
- 401/403 → auth_error(不重试,由调用方禁用账号并通知)
- 其他 4xx → 直接失败(配置问题,不重试)
- 网络异常 / 5xx → 重试 retry_count 次,间隔 2s
- JSON 解析/路径提取失败 → 直接失败
"""
url, headers, body = _build_request(platform, api_key)
method = platform["method"]
attempts = retry_count + 1
for attempt in range(attempts):
try:
if method == "GET":
resp = requests.get(url, headers=headers, timeout=timeout)
else:
resp = requests.post(url, headers=headers, data=body, timeout=timeout)
except requests.RequestException as exc:
logger.warning("请求异常(%s/%s) %s: %s", attempt + 1, attempts, url, exc)
if attempt < attempts - 1:
time.sleep(2)
continue
if resp.status_code in (401, 403):
return FetchResult(ok=False, auth_error=True, status_code=resp.status_code,
error=f"HTTP {resp.status_code}(认证失败)")
if resp.status_code >= 500:
logger.warning("服务端错误(%s/%s) %s status=%s", attempt + 1, attempts, url, resp.status_code)
if attempt < attempts - 1:
time.sleep(2)
continue
if resp.status_code >= 400:
return FetchResult(ok=False, status_code=resp.status_code,
error=f"HTTP {resp.status_code}: {resp.text[:120]}")
try:
data = resp.json()
except ValueError:
return FetchResult(ok=False, status_code=resp.status_code,
error="响应不是合法 JSON")
try:
balance = extract_balance(data, platform["balance_path"])
except ValueError as exc:
return FetchResult(ok=False, status_code=resp.status_code, error=str(exc))
return FetchResult(ok=True, balance=balance, status_code=resp.status_code)
return FetchResult(ok=False, error=f"网络错误,重试 {retry_count} 次后仍失败")