import time from typing import Optional import requests as req_lib class CookiePool: STATUS_INVALID = "2" STATUS_MUTED = "3" def __init__(self, session_api: str, app_id: str, secret: str, *, retries: int = 3, auth_token: str = ""): self.session_api = session_api.rstrip("/") self.app_id = app_id self.secret = secret self.retries = max(1, int(retries)) self.auth_token = auth_token.strip() self._cache: dict[str, dict] = {} def _headers(self) -> dict: h = {"Accept": "application/json"} if self.auth_token: token = self.auth_token if not token.lower().startswith("bearer "): token = f"Bearer {token}" h["Authorization"] = token return h @classmethod def from_config(cls, cfg: dict) -> "CookiePool": cp = cfg.get("cookie_pool") or {} return cls( session_api=cp.get("session_api") or f"{cfg.get('api_base', '').rstrip('/')}/api/third", app_id=cp.get("app_id") or "", secret=cp.get("secret") or "", retries=int(cp.get("retries") or 3), auth_token=cp.get("auth_token") or "", ) def get_cookie(self, platform_id: str | int, *, retries: int | None = None) -> Optional[dict]: pid = str(platform_id) if not self.app_id or not self.secret: print(" [cookie-pool] 缺少 app_id/secret, 请检查配置 cookie_pool") return None n = self.retries if retries is None else max(1, int(retries)) url = f"{self.session_api}/getOneSpiderSession?platform_id={pid}&app_id={self.app_id}&secret={self.secret}" for attempt in range(n): try: resp = req_lib.get(url, headers=self._headers(), timeout=30) data = resp.json() session = data.get("data") if not session or session == []: print(f" [cookie-pool] platform_id={pid} 无可用 session: {data.get('msg') or data}") if attempt + 1 < n: time.sleep(2) continue return None cookie = session.get("cookie", "") sid = session.get("id", "") if not cookie: print(f" [cookie-pool] platform_id={pid} 无可用 cookie: {data.get('msg') or data}") if attempt + 1 < n: time.sleep(2) continue return None print(f" [cookie-pool] platform_id={pid} 获取成功 id={sid} cookie={str(cookie)[:36]}...") return {"cookie": cookie, "id": sid} except Exception as e: print(f" [cookie-pool] platform_id={pid} 获取异常({attempt+1}/{n}): {e}") if attempt + 1 < n: time.sleep(2) return None def get_for_internal(self, internal: str, *, retries: int | None = None) -> Optional[dict]: from config import COOKIE_PLATFORM_IDS pid = COOKIE_PLATFORM_IDS.get(internal) if not pid: return None return self.get_cookie(pid, retries=retries) def update_session(self, session_id, reload_time: str = "", status: str = "2") -> bool: if session_id in (None, "", "-1", -1): return False if not self.app_id or not self.secret: return False url = f"{self.session_api}/updateSpiderSession?app_id={self.app_id}&secret={self.secret}" try: resp = req_lib.post(url, headers=self._headers(), json={"id": session_id, "status": str(status), "reload_time": reload_time or ""}, timeout=10) print(f" [cookie-pool] update id={session_id} status={status} -> {resp.text[:120]}") return True except Exception as e: print(f" [cookie-pool] update 失败 id={session_id}: {e}") return False def invalidate(self, session_id, reason: str = "invalid"): status = self.STATUS_MUTED if reason in ("muted", "quota", "rate_limit") else self.STATUS_INVALID return self.update_session(session_id, "", status) def parse_cookie_string(cookie_str: str) -> dict[str, str]: if not cookie_str: return {} s = cookie_str.strip() if s.startswith("{"): return {} pairs = {} for part in s.split(";"): part = part.strip() if "=" in part: k, v = part.split("=", 1) pairs[k.strip()] = v.strip() return pairs def cookie_str_to_playwright(cookie_str: str, domain: str | list[str]) -> list[dict]: cookies = [] domains = [domain] if isinstance(domain, str) else domain for name, value in parse_cookie_string(cookie_str).items(): for d in domains: cookies.append({"name": name, "value": value, "domain": d, "path": "/"}) return cookies