task_worker_v2_portable/cookie_pool.py

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