"""Upstream HTTP client — ported from monitor_ai98pro_group_rates.py.""" from __future__ import annotations import json from decimal import Decimal, InvalidOperation from typing import Any, Optional from urllib.parse import urljoin import httpx class UpstreamError(RuntimeError): pass def _find_token(value: Any) -> str: if isinstance(value, str) and value.count(".") >= 2: return value if isinstance(value, dict): for key in ("token", "access_token", "accessToken", "jwt", "auth_token", "authToken"): candidate = value.get(key) if isinstance(candidate, str) and candidate: return candidate for key in ("data", "result", "user", "session"): tok = _find_token(value.get(key)) if tok: return tok return "" def _find_user_id(value: Any) -> str: if isinstance(value, dict): for key in ("id", "user_id", "userId"): candidate = value.get(key) if candidate is not None: return str(candidate) for key in ("data", "result", "user", "session"): user_id = _find_user_id(value.get(key)) if user_id: return user_id return "" def _unwrap_list(value: Any) -> Optional[list[dict[str, Any]]]: def _normalize(lst: list) -> list[dict[str, Any]]: out = [] for i in lst: if isinstance(i, dict): out.append(i) elif isinstance(i, str): out.append({"id": i, "name": i}) return out if isinstance(value, list): return _normalize(value) if isinstance(value, dict): for key in ("data", "items", "groups", "available_groups", "availableGroups"): nested = value.get(key) if isinstance(nested, list): return _normalize(nested) elif isinstance(nested, dict): # Handle /api/user/self/groups where data is a dict of group_name -> { desc, ratio } out = [] for k in nested.keys(): out.append({"id": k, "name": k}) return out return None def _decimal_str(value: Any) -> str: if value is None or value == "": return "" try: d = Decimal(str(value)) except (InvalidOperation, ValueError): return str(value) n = d.normalize() if n == n.to_integral(): return str(n.quantize(Decimal("1"))) return format(n, "f") def _group_id(group: dict[str, Any]) -> str: for key in ("id", "group_id", "groupId"): v = group.get(key) if v is not None: return str(v) name = str(group.get("name") or group.get("group_name") or "") platform = str(group.get("platform") or "") return f"{platform}:{name}" def _rate_from_group(group: dict[str, Any]) -> str: for key in ( "user_rate_multiplier", "userRateMultiplier", "effective_rate_multiplier", "effectiveRateMultiplier", "rate_multiplier", "rateMultiplier", ): r = _decimal_str(group.get(key)) if r: return r return "" def _extract_rates_map(raw: Any) -> dict[str, str]: if raw is None: return {} # Handle one-api/new-api /api/option response where GroupRatio is in a list of options if isinstance(raw, dict) and isinstance(raw.get("data"), list): for item in raw["data"]: if isinstance(item, dict) and item.get("key") == "GroupRatio": val = item.get("value") if isinstance(val, str): try: import json parsed = json.loads(val) if isinstance(parsed, dict): result: dict[str, str] = {} for k, v in parsed.items(): r = _decimal_str(v) if r: result[str(k)] = r return result except Exception: pass elif isinstance(val, dict): # In case it's returned as dict directly result = {} for k, v in val.items(): r = _decimal_str(v) if r: result[str(k)] = r return result if isinstance(raw, dict): candidates = raw for key in ("data", "rates", "group_rates", "groupRates", "GroupRatio"): nested = raw.get(key) if isinstance(nested, dict): candidates = nested break elif isinstance(nested, str) and key == "GroupRatio": # Handle GroupRatio as a JSON string try: import json parsed = json.loads(nested) if isinstance(parsed, dict): candidates = parsed break except Exception: pass result: dict[str, str] = {} for k, v in candidates.items(): if isinstance(v, dict): r = _decimal_str( v.get("rate_multiplier") or v.get("rateMultiplier") or v.get("user_rate_multiplier") or v.get("userRateMultiplier") or v.get("ratio") ) else: r = _decimal_str(v) if r: result[str(k)] = r return result if isinstance(raw, list): result = {} for item in raw: if not isinstance(item, dict): continue gid = _group_id(item) rate = _rate_from_group(item) if gid and rate: result[gid] = rate return result return {} def build_snapshot(upstream_id: int, base_url: str, api_prefix: str, groups: list[dict[str, Any]], raw_rates: Any) -> dict[str, Any]: from datetime import datetime, timezone override_rates = _extract_rates_map(raw_rates) entries: dict[str, dict[str, Any]] = {} for g in groups: gid = _group_id(g) default_rate = _rate_from_group(g) effective_rate = override_rates.get(gid, default_rate) entries[gid] = { "group_id": gid, "group_name": g.get("name") or g.get("group_name") or "", "platform": g.get("platform") or "", "rate": effective_rate, "default_rate": default_rate, "override_rate": override_rates.get(gid, ""), } return { "upstream_id": upstream_id, "base_url": base_url.rstrip("/"), "api_prefix": api_prefix, "captured_at": datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds"), "groups": entries, } class UpstreamClient: """Sync HTTP client that handles all auth types.""" def __init__( self, base_url: str, api_prefix: str, auth_type: str, auth_config: dict[str, Any], timeout: float = 30.0, ) -> None: self.base_url = base_url.rstrip("/") self.api_prefix = api_prefix.strip("/") self.auth_type = auth_type self.auth_config = auth_config self.timeout = timeout self._token: str = "" self._cookies: dict[str, str] = {} self._new_api_user: str = "" def _url(self, path: str) -> str: prefix = f"/{self.api_prefix}" if self.api_prefix else "" return f"{self.base_url}{prefix}/{path.lstrip('/')}" def _headers(self, auth: bool = True) -> dict[str, str]: headers: dict[str, str] = { "Accept": "application/json", "User-Agent": "SmartUp/1.0", } if not auth: return headers if self.auth_type == "bearer": token = self.auth_config.get("token", "") if token: headers["Authorization"] = f"Bearer {token}" elif self.auth_type == "api_key": key = self.auth_config.get("key", "") header = self.auth_config.get("header", "Authorization") if key: headers[header] = key elif self.auth_type == "login_password" and self._token: headers["Authorization"] = f"Bearer {self._token}" if self.auth_type == "login_password" and self._new_api_user: headers["New-Api-User"] = self._new_api_user return headers def _request(self, method: str, path: str, body: Any = None, auth: bool = True) -> Any: url = self._url(path) with httpx.Client(timeout=self.timeout) as client: if body is not None: resp = client.request( method, url, json=body, headers=self._headers(auth), cookies=self._cookies, ) else: resp = client.request( method, url, headers=self._headers(auth), cookies=self._cookies, ) self._cookies.update(dict(resp.cookies)) resp.raise_for_status() ct = resp.headers.get("content-type", "") if not resp.content: return None text = resp.text if "application/json" not in ct and text.lstrip().startswith("<"): raise UpstreamError(f"{method} {path} returned HTML, not JSON") return resp.json() def login(self) -> None: if self.auth_type != "login_password": return email = self.auth_config.get("email", "") password = self.auth_config.get("password", "") login_path = self.auth_config.get("login_path", "/auth/login") username_field = self.auth_config.get("username_field", "email") if not email or not password: raise UpstreamError("login_password auth requires email and password in auth_config") resp = self._request("POST", login_path, {username_field: email, "password": password}, auth=False) token = _find_token(resp) if token: self._token = token return if self._cookies: self._new_api_user = self.auth_config.get("new_api_user", "") or _find_user_id(resp) return raise UpstreamError("login succeeded but no token or session cookie found in response") def get_available_groups(self, endpoint: str) -> list[dict[str, Any]]: resp = self._request("GET", endpoint) groups = _unwrap_list(resp) if groups is None: raise UpstreamError(f"{endpoint} did not return a list") return groups def get_group_rates(self, endpoint: str) -> Any: return self._request("GET", endpoint)