import json import logging import time from datetime import datetime, timezone from typing import Any, Dict, List, Tuple from curl_cffi import requests as cffi_requests logger = logging.getLogger(__name__) class Sub2APIClient: def __init__(self, api_url: str, api_key: str): self.api_url = api_url.rstrip("/") self.headers = { "Content-Type": "application/json", "x-api-key": api_key, } self.request_kwargs = { "timeout": 15, "impersonate": "chrome110", } def _handle_response( self, response: cffi_requests.Response, success_codes: Tuple[int, ...] = (200, 201, 204), ) -> Tuple[bool, Any]: if response.status_code in success_codes: try: return True, response.json() if response.text else {} except ValueError: return True, response.text error_msg = f"HTTP {response.status_code}" try: detail = response.json() if isinstance(detail, dict): error_msg = detail.get("message", error_msg) except Exception: error_msg = f"{error_msg} - {response.text[:200]}" return False, error_msg def _get_push_settings(self) -> Dict[str, Any]: try: import utils.config as cfg except ImportError: cfg = None def as_int(value: Any, default: int, minimum: int) -> int: try: return max(minimum, int(value)) except (TypeError, ValueError): return default def as_float(value: Any, default: float, minimum: float) -> float: try: return max(minimum, float(value)) except (TypeError, ValueError): return default raw_group_ids = getattr(cfg, "SUB2API_ACCOUNT_GROUP_IDS", []) if cfg else [] if isinstance(raw_group_ids, list): group_ids = [int(item) for item in raw_group_ids if str(item).strip().isdigit()] else: group_ids = [int(item.strip()) for item in str(raw_group_ids or "").split(",") if item.strip().isdigit()] return { "concurrency": as_int(getattr(cfg, "SUB2API_ACCOUNT_CONCURRENCY", 10) if cfg else 10, 10, 1), "load_factor": as_int(getattr(cfg, "SUB2API_ACCOUNT_LOAD_FACTOR", 10) if cfg else 10, 10, 1), "priority": as_int(getattr(cfg, "SUB2API_ACCOUNT_PRIORITY", 1) if cfg else 1, 1, 1), "rate_multiplier": as_float(getattr(cfg, "SUB2API_ACCOUNT_RATE_MULTIPLIER", 1.0) if cfg else 1.0, 1.0, 0.0), "group_ids": group_ids, "enable_ws": bool(getattr(cfg, "SUB2API_ENABLE_WS_MODE", True) if cfg else True), } def _build_account_extra(self, settings: Dict[str, Any]) -> Dict[str, Any]: extra = {"load_factor": settings["load_factor"]} if settings["enable_ws"]: extra["openai_oauth_responses_websockets_v2_enabled"] = True extra["openai_oauth_responses_websockets_v2_mode"] = "passthrough" return extra def _refresh_created_account(self, account_id: str) -> None: if not account_id: return refresh_urls = [ f"{self.api_url}/api/v1/admin/accounts/{account_id}/refresh", f"{self.api_url}/api/v1/admin/openai/accounts/{account_id}/refresh", ] for refresh_url in refresh_urls: try: response = cffi_requests.post( refresh_url, json={}, headers=self.headers, timeout=30, impersonate="chrome110", proxies=None, ) if response.status_code in (200, 201, 204): logger.info("Sub2API account refresh succeeded (ID: %s)", account_id) return except Exception as exc: logger.warning("Sub2API account refresh failed via %s: %s", refresh_url, exc) logger.warning("Sub2API account refresh did not succeed for %s", account_id) def _import_account(self, token_data: Dict[str, Any], settings: Dict[str, Any]) -> Tuple[bool, str]: url = f"{self.api_url}/api/v1/admin/accounts/data" exported_at = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") extra = self._build_account_extra(settings) account_item = { "name": token_data.get("email", "unknown"), "platform": "openai", "type": "oauth", "credentials": { "access_token": token_data.get("access_token", ""), "chatgpt_account_id": token_data.get("account_id", ""), "client_id": token_data.get("client_id", ""), "expires_at": int(time.time() + 864000), "expires_in": 863999, "model_mapping": { "gpt-4o": "gpt-4o", "gpt-4": "gpt-4", "gpt-3.5-turbo": "gpt-3.5-turbo", }, "organization_id": token_data.get("workspace_id", ""), "refresh_token": token_data.get("refresh_token", ""), }, "extra": extra, "concurrency": settings["concurrency"], "priority": settings["priority"], "rate_multiplier": settings["rate_multiplier"], "auto_pause_on_expired": True, } if settings["group_ids"]: account_item["group_ids"] = settings["group_ids"] payload = { "data": { "type": "sub2api-data", "version": 1, "exported_at": exported_at, "proxies": [], "accounts": [account_item], }, "skip_default_group_bind": not bool(settings["group_ids"]), } try: headers = self.headers.copy() headers["Idempotency-Key"] = f"import-{int(time.time())}" response = cffi_requests.post( url, json=payload, headers=headers, timeout=30, impersonate="chrome110", proxies=None, ) ok, result = self._handle_response(response, success_codes=(200, 201)) if ok: return True, "Sub2API account import succeeded" return False, str(result) except Exception as exc: return False, f"Network request failed: {exc}" def get_accounts(self, page: int = 1, page_size: int = 50) -> Tuple[bool, Any]: url = f"{self.api_url}/api/v1/admin/accounts" params = { "page": page, "page_size": page_size, } try: response = cffi_requests.get(url, headers=self.headers, params=params, **self.request_kwargs) return self._handle_response(response) except Exception as exc: logger.error("Get Sub2API accounts failed: %s", exc) return False, str(exc) def get_all_accounts(self, page_size: int = 100) -> Tuple[bool, Any]: all_items: List[dict] = [] page = 1 while True: ok, data = self.get_accounts(page=page, page_size=page_size) if not ok: if page == 1: return False, data logger.warning( "Sub2API pagination failed on page %s; continue with %s fetched accounts", page, len(all_items), ) break inner = data.get("data", {}) if isinstance(data, dict) else {} items = inner.get("items", []) if not items: break all_items.extend(items) total = inner.get("total", 0) if len(all_items) >= total: break page += 1 logger.info("Fetched %s Sub2API accounts across paginated results", len(all_items)) return True, all_items def add_account(self, token_data: Dict[str, Any]) -> Tuple[bool, str]: settings = self._get_push_settings() refresh_token = token_data.get("refresh_token", "") if not refresh_token: return self._import_account(token_data, settings) url = f"{self.api_url}/api/v1/admin/accounts" payload = { "name": token_data.get("email", "unknown")[:64], "platform": "openai", "type": "oauth", "credentials": {"refresh_token": refresh_token}, "concurrency": settings["concurrency"], "priority": settings["priority"], "rate_multiplier": settings["rate_multiplier"], "extra": self._build_account_extra(settings), } if settings["group_ids"]: payload["group_ids"] = settings["group_ids"] try: response = cffi_requests.post( url, json=payload, headers=self.headers, timeout=30, impersonate="chrome110", proxies=None, ) ok, result = self._handle_response(response, success_codes=(200, 201)) if not ok: logger.warning("Sub2API direct create failed, falling back to import endpoint: %s", result) return self._import_account(token_data, settings) account_id = result.get("data", {}).get("id") if isinstance(result, dict) else None if account_id: self._refresh_created_account(str(account_id)) return True, "Sub2API account created successfully" except Exception as exc: logger.warning("Sub2API direct create raised an exception, falling back to import: %s", exc) return self._import_account(token_data, settings) def update_account(self, account_id: str, update_data: Dict[str, Any]) -> Tuple[bool, Any]: url = f"{self.api_url}/api/v1/admin/accounts/{account_id}" try: response = cffi_requests.put(url, json=update_data, headers=self.headers, **self.request_kwargs) return self._handle_response(response) except Exception as exc: logger.error("Update Sub2API account %s failed: %s", account_id, exc) return False, str(exc) def set_account_status(self, account_id: str, disabled: bool) -> bool: url = f"{self.api_url}/api/v1/admin/accounts/{account_id}" status_val = "inactive" if disabled else "active" payload = {"status": status_val} try: response = cffi_requests.patch(url, json=payload, headers=self.headers, **self.request_kwargs) if response.status_code in (200, 201, 204): return True response = cffi_requests.put(url, json=payload, headers=self.headers, **self.request_kwargs) return response.status_code in (200, 201, 204) except Exception as exc: logger.error("Set Sub2API account %s status failed: %s", account_id, exc) return False def delete_account(self, account_id: str) -> Tuple[bool, Any]: url = f"{self.api_url}/api/v1/admin/accounts/{account_id}" try: response = cffi_requests.delete(url, headers=self.headers, **self.request_kwargs) return self._handle_response(response, success_codes=(200, 204)) except Exception as exc: logger.error(f"删除账号 {account_id} 失败: {exc}") return False, str(exc) def refresh_account(self, account_id: str) -> Tuple[bool, Any]: url = f"{self.api_url}/api/v1/admin/accounts/{account_id}/refresh" try: response = cffi_requests.post(url, headers=self.headers, json={}, **self.request_kwargs) return self._handle_response(response) except Exception as exc: logger.error(f"刷新账号 {account_id} 失败: {exc}") return False, str(exc) def test_account(self, account_id: int) -> Tuple[str, str]: url = f"{self.api_url}/api/v1/admin/accounts/{account_id}/test" try: response = cffi_requests.post( url, headers=self.headers, json={}, timeout=60, impersonate="chrome110", ) if response.status_code != 200: logger.warning("Sub2API test_account %s returned HTTP %s; keep current state", account_id, response.status_code) return "ok", f"HTTP {response.status_code}, skipped" for line in response.text.splitlines(): line = line.strip() if not line.startswith("data:"): continue raw = line[5:].strip() if not raw or raw == "[DONE]": continue try: event = json.loads(raw) except Exception: continue event_type = event.get("type", "") if event_type == "test_complete": if event.get("success"): return "ok", "test completed" err = str(event.get("error") or event.get("text") or "") return _classify_sse_error(err) if event_type == "error": err = str(event.get("error") or event.get("text") or "") return _classify_sse_error(err) logger.warning("Sub2API test_account %s did not emit a terminal SSE event; keep current state", account_id) return "ok", "no terminal SSE event, skipped" except Exception as exc: logger.warning("Sub2API test_account %s failed: %s", account_id, exc) return "ok", f"test error, skipped: {str(exc)[:120]}" def test_connection(self) -> Tuple[bool, str]: url = f"{self.api_url}/api/v1/admin/accounts/data" try: kwargs = self.request_kwargs.copy() kwargs["timeout"] = 10 response = cffi_requests.get(url, headers=self.headers, **kwargs) if response.status_code in (200, 201, 204, 405): return True, "Sub2API connection test succeeded. The API key is valid." if response.status_code == 401: return False, "Connected, but the API key is invalid (401 Unauthorized)." if response.status_code == 403: return False, "Connected, but the API key does not have enough permission (403 Forbidden)." return False, f"Unexpected server status code: {response.status_code}" except cffi_requests.exceptions.ConnectionError as exc: return False, f"Could not connect to the Sub2API server: {exc}" except cffi_requests.exceptions.Timeout: return False, "连接超时，请检查网络配置或服务器状态" except Exception as exc: return False, f"连接测试失败: {str(exc)}" def _classify_sse_error(err_text: str) -> Tuple[str, str]: text = err_text.lower() if any(keyword in text for keyword in ("429", "rate_limit", "rate limit", "too many request")): return "quota", f"quota limited: {err_text[:120]}" if err_text.strip(): return "dead", f"test failed: {err_text[:120]}" return "ok", "empty SSE error, skipped"