agento3/proxy_pool.py

"""
异步代理池 v4 — Cookie 感知 + 挑战分离 + 真实统计

v4 vs v3:
- 每个代理维护 cookie jar（cf_clearance 等），成功请求后自动保存
- CF 挑战不再标记为代理死亡（mark_challenged ≠ mark_failure）
- 真实统计：维护原子计数器，不造假数据
- 更多 TLS 指纹：加入 safari15_5/safari17_0
"""
import asyncio
import logging
import os
import random
import time
from email.utils import parsedate_to_datetime
from threading import Lock
from typing import Optional

from curl_cffi.requests import AsyncSession

logger = logging.getLogger("ao3-proxy-pool")

OPTIMAL_MIN_PORT = 13500
OPTIMAL_MAX_PORT = 14499
WORKING_PROXIES_FILE = "/dev/shm/working_proxies.txt"

FAST_POOL_SIZE = 50

# 被动检查：只对快池做轻量采样
SAMPLE_INTERVAL = 60
SAMPLE_BATCH_SIZE = 10
SAMPLE_TIMEOUT = 5
FAST_POOL_REFRESH_INTERVAL = 10

# TLS fingerprints — proven against AO3 Cloudflare
# safari15_5/17_0 have highest CF bypass rate per testing
BROWSER_IMPS = ["safari15_5", "safari17_0", "chrome123", "chrome124"]

WARM_THRESHOLD_S = 3.0


def _parse_cookie_expires(set_cookie: str) -> float:
    """Extract expiry timestamp from Set-Cookie header. Returns 0 if session cookie."""
    for part in set_cookie.split(";"):
        part = part.strip()
        if part.lower().startswith("expires="):
            try:
                dt = parsedate_to_datetime(part[8:])
                if dt:
                    return dt.timestamp()
            except Exception:
                pass
        elif part.lower().startswith("max-age="):
            try:
                return time.time() + int(part[9:])
            except Exception:
                pass
    return 0  # session cookie


def load_proxies_from_file() -> list[str]:
    """Load proxy list from known-working file, or full list with port filtering."""
    if os.path.exists(WORKING_PROXIES_FILE) and os.path.getsize(WORKING_PROXIES_FILE) > 0:
        with open(WORKING_PROXIES_FILE) as f:
            proxies = [l.strip() for l in f if l.strip() and ":" in l]
        if proxies:
            logger.info(f"Loaded {len(proxies)} working proxies from {WORKING_PROXIES_FILE}")
            filtered = [p for p in proxies if ":" in p and OPTIMAL_MIN_PORT <= int(p.split(":")[-1]) <= OPTIMAL_MAX_PORT]
            if len(filtered) >= 10:
                return filtered
            return proxies
    proxy_file = "/home/ubuntu/proxy.txt"
    proxies = []
    with open(proxy_file) as f:
        for line in f:
            line = line.strip()
            if not line or ":" not in line:
                continue
            if "|" in line:
                line = line.split("|")[-1].strip()
            proxies.append(line)
    filtered = [p for p in proxies if ":" in p and OPTIMAL_MIN_PORT <= int(p.split(":")[-1]) <= OPTIMAL_MAX_PORT]
    return filtered if len(filtered) >= 10 else proxies


class ProxySession:
    """A single proxy with its own AsyncSession, cookie jar, and health state."""

    __slots__ = (
        "host_port", "host", "port_str", "port",
        "session", "impersonate", "alive",
        "consecutive_failures", "ban_until", "last_used",
        "avg_response_time", "weight", "requests_handled",
        "last_sample", "sample_passed",
        # v4: cookie jar
        "_cookies", "_cookie_expires", "_lock",
    )

    def __init__(self, host_port: str):
        self.host_port = host_port
        self.host, self.port_str = host_port.split(":")
        self.port = int(self.port_str)
        self.session: Optional[AsyncSession] = None
        self.impersonate = random.choice(BROWSER_IMPS)
        self.alive = True
        self.consecutive_failures = 0
        self.ban_until = 0.0
        self.last_used = 0.0
        self.avg_response_time = 1.0
        self.weight = 1.0
        self.requests_handled = 0
        self.last_sample = 0.0
        self.sample_passed = True
        # v4: per-proxy cookie jar (cf_clearance, etc.)
        self._cookies: dict[str, str] = {}
        self._cookie_expires: dict[str, float] = {}
        self._lock = Lock()

    # ─── Cookie management ────────────────────────────────────────────

    def save_cookies(self, headers: dict) -> int:
        """Extract Set-Cookie from response headers. Returns count saved."""
        saved = 0
        set_cookie = headers.get("Set-Cookie", "")
        if not set_cookie:
            # Some servers use set-cookie (lowercase) in HTTP/2
            set_cookie = headers.get("set-cookie", "")
        if not set_cookie:
            return 0

        now = time.time()
        with self._lock:
            for part in set_cookie.split(","):
                # Handle comma-separated cookies (ugh)
                part = part.strip()
                if "=" not in part:
                    continue
                name, _, rest = part.partition("=")
                value = rest.split(";")[0].strip() if ";" in rest else rest.strip()
                name = name.strip()
                if not name:
                    continue
                self._cookies[name] = value
                expires = _parse_cookie_expires(part)
                if expires > 0:
                    self._cookie_expires[name] = expires
                saved += 1

        # Purge expired cookies
        self._purge_expired(now)
        if saved:
            logger.debug(f"Saved {saved} cookies for {self.host_port} (keys: {list(self._cookies.keys())})")
        return saved

    def get_cookie_header(self) -> str:
        """Get Cookie header string for this proxy. Returns '' if no cookies."""
        now = time.time()
        with self._lock:
            self._purge_expired(now)
            if not self._cookies:
                return ""
            return "; ".join(f"{k}={v}" for k, v in self._cookies.items())

    def _purge_expired(self, now: float):
        """Remove expired cookies."""
        expired = [k for k, exp in self._cookie_expires.items() if 0 < exp < now]
        for k in expired:
            self._cookies.pop(k, None)
            self._cookie_expires.pop(k, None)

    # ─── Session management ───────────────────────────────────────────

    async def get_session(self) -> AsyncSession:
        if self.session is None:
            self.session = AsyncSession()
            self.session.impersonate = self.impersonate
            self.session.timeout = 15
            self.session.proxies = {
                "http": f"http://{self.host_port}",
                "https": f"http://{self.host_port}",
            }
            # Default headers that don't change per-request (safe to set on session)
            self.session.headers.update({
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                              "(KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36",
                "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
                "Accept-Language": "en-US,en;q=0.9",
                "Accept-Encoding": "gzip, deflate, br",
            })
        return self.session

    async def close(self):
        if self.session:
            try:
                await self.session.close()
            except Exception:
                pass
            self.session = None

    # ─── Health tracking ──────────────────────────────────────────────

    def update_weight(self):
        self.weight = 1.0 / max(self.avg_response_time, 0.1)

    def mark_success(self, response_time: float):
        self.alive = True
        self.consecutive_failures = 0
        self.ban_until = 0.0
        self.requests_handled += 1
        self.avg_response_time = self.avg_response_time * 0.7 + response_time * 0.3
        self.update_weight()

    def mark_failure(self):
        """Proxy-level failure (connection error, timeout, etc.) — exponential backoff."""
        self.requests_handled += 1
        self.consecutive_failures += 1
        backoff = min(5 * (3 ** (self.consecutive_failures - 1)), 300)
        self.ban_until = time.time() + backoff
        if self.consecutive_failures >= 3:
            self.alive = False
        self.sample_passed = False

    def mark_challenged(self):
        """CF challenge detected — proxy is alive but needs cookie. NO backoff."""
        # Don't increment consecutive_failures — challenge is not proxy's fault
        # Don't set ban_until — proxy may work with cookies
        self.requests_handled += 1
        # Only mark as not-sample-passed so it won't be fast-pool priority
        self.sample_passed = False

    def mark_sample(self, passed: bool, response_time: float = 0):
        """Lightweight periodic check result."""
        self.last_sample = time.time()
        self.sample_passed = passed
        if passed and not self.alive:
            self.alive = True
            self.consecutive_failures = max(0, self.consecutive_failures - 1)
            self.avg_response_time = self.avg_response_time * 0.5 + response_time * 0.5
            self.update_weight()

    @property
    def is_available(self) -> bool:
        return self.alive and time.time() > self.ban_until

    def __repr__(self):
        nc = len(self._cookies)
        return f"PS({self.host_port}, alive={self.alive}, {self.avg_response_time:.1f}s, cookies={nc})"


class AsyncProxyPool:
    """Tiered async proxy pool with cookie-aware session management."""

    def __init__(self):
        self._proxies: list[ProxySession] = []
        self._fast_pool: list[ProxySession] = []
        self._fast_pool_updated = 0.0
        self._sample_task: Optional[asyncio.Task] = None
        # v4: atomic counters for real stats (no more fake data)
        self._stats_lock = Lock()
        self._stats = {"alive": 0, "dead": 0, "banned": 0, "challenged": 0}
        self._stats_cache = {}
        self._stats_cache_ts = 0.0
        self._load_proxies()
        self._start_sampling()

    def _load_proxies(self):
        proxies = load_proxies_from_file()
        self._proxies = [ProxySession(hp) for hp in proxies]
        self._refresh_fast_pool()
        self._recompute_stats()
        logger.info(f"ProxyPool v4 ready: {len(self._proxies)} proxies, fast={len(self._fast_pool)}")

    def _start_sampling(self):
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
        if self._sample_task is None or self._sample_task.done():
            self._sample_task = asyncio.create_task(self._sampling_loop())

    def _refresh_fast_pool(self):
        """Select fastest N proxies, seeded immediately from all alive."""
        alive = [p for p in self._proxies if p.is_available]
        sampled = [p for p in alive if p.sample_passed]
        unsampled = [p for p in alive if not p.sample_passed]
        sampled.sort(key=lambda p: p.avg_response_time)
        unsampled.sort(key=lambda p: p.avg_response_time)
        combined = sampled + unsampled
        self._fast_pool = combined[:FAST_POOL_SIZE]
        self._fast_pool_updated = time.time()

    def _recompute_stats(self):
        """Accurate stats — iterate all proxies (fast, ~726 items)."""
        alive = 0
        dead = 0
        banned = 0
        for p in self._proxies:
            if p.alive:
                alive += 1
                if not p.is_available:
                    banned += 1
            else:
                dead += 1
        with self._stats_lock:
            self._stats = {"alive": alive, "dead": dead, "banned": banned}

    # ─── Proxy selection ──────────────────────────────────────────────

    def get_fast_proxy(self) -> Optional[ProxySession]:
        """Fast proxy for interactive requests (login/register/POST)."""
        now = time.time()
        if now - self._fast_pool_updated > FAST_POOL_REFRESH_INTERVAL:
            self._refresh_fast_pool()
        if not self._fast_pool:
            return self.get_proxy()
        total = sum(p.weight for p in self._fast_pool)
        if total <= 0:
            return random.choice(self._fast_pool)
        r = random.uniform(0, total)
        cum = 0
        for p in self._fast_pool:
            cum += p.weight
            if r <= cum:
                return p
        return random.choice(self._fast_pool)

    def get_proxy(self) -> Optional[ProxySession]:
        """Weighted random from all available proxies."""
        available = [p for p in self._proxies if p.is_available]
        if not available:
            # Fallback: use proxies with fewer than 10 failures
            available = [p for p in self._proxies if p.consecutive_failures < 10]
        if not available:
            return None
        total = sum(p.weight for p in available)
        if total <= 0:
            return random.choice(available)
        r = random.uniform(0, total)
        cum = 0
        for p in available:
            cum += p.weight
            if r <= cum:
                return p
        return random.choice(available)

    def get_proxy_with_cookies(self) -> Optional[ProxySession]:
        """Get a proxy that has cookies saved (cf_clearance). Fallback to any available."""
        available = [p for p in self._proxies if p.is_available]
        with_cookies = [p for p in available if p._cookies]
        if with_cookies:
            total = sum(p.weight for p in with_cookies)
            if total > 0:
                r = random.uniform(0, total)
                cum = 0
                for p in with_cookies:
                    cum += p.weight
                    if r <= cum:
                        return p
            return random.choice(with_cookies)
        return self.get_proxy()

    # ─── Sampling loop ────────────────────────────────────────────────

    async def _sampling_loop(self):
        """Lightweight sampling — only test fast-pool proxies."""
        logger.info("Sampling loop started (lightweight, fast-pool only)")
        while True:
            try:
                await asyncio.sleep(SAMPLE_INTERVAL)
                pool = self._fast_pool[:] if self._fast_pool else self._proxies[:50]
                if not pool:
                    continue
                alive_cnt = 0
                dead_cnt = 0
                for i in range(0, len(pool), SAMPLE_BATCH_SIZE):
                    batch = pool[i:i + SAMPLE_BATCH_SIZE]
                    checks = [self._check_single(p) for p in batch]
                    results = await asyncio.gather(*checks, return_exceptions=True)
                    for p, r in zip(batch, results):
                        if isinstance(r, Exception):
                            p.mark_sample(False)
                            dead_cnt += 1
                        elif r[0]:
                            p.mark_sample(True, r[1])
                            alive_cnt += 1
                        else:
                            p.mark_sample(False)
                            dead_cnt += 1
                self._refresh_fast_pool()
                self._recompute_stats()
                total_avg = sum(p.avg_response_time for p in pool if p.requests_handled > 0) / max(alive_cnt, 1)
                logger.debug(f"Sample: {alive_cnt} alive, {dead_cnt} dead, fast={len(self._fast_pool)}, avg={total_avg*1000:.0f}ms")
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Sample error: {e}")

    async def _check_single(self, proxy: ProxySession) -> tuple[bool, float]:
        """Quick single-proxy check against AO3."""
        start = time.time()
        try:
            s = await proxy.get_session()
            resp = await s.head("https://archiveofourown.org/", timeout=SAMPLE_TIMEOUT)
            return (200 <= resp.status_code < 500, time.time() - start)
        except Exception:
            return (False, time.time() - start)

    # ─── Lifecycle ────────────────────────────────────────────────────

    async def close_all(self):
        if self._sample_task and not self._sample_task.done():
            self._sample_task.cancel()
            try:
                await self._sample_task
            except asyncio.CancelledError:
                pass
        for p in self._proxies:
            await p.close()
        logger.info("All sessions closed")

    # ─── Stats ────────────────────────────────────────────────────────

    def get_stats(self) -> dict:
        """Accurate stats with caching (1s throttle to avoid iteration on every call)."""
        now = time.time()
        if self._stats_cache and now - self._stats_cache_ts < 1.0:
            return self._stats_cache

        self._recompute_stats()
        total = len(self._proxies)

        # Compute avg response from sampled proxies
        sampled = [p.avg_response_time for p in self._proxies if p.requests_handled > 0]
        avg_speed = sum(sampled) / max(len(sampled), 1)

        # Count proxies with cookies
        with_cookies = sum(1 for p in self._proxies if p._cookies)

        result = {
            "total": total,
            "alive": self._stats["alive"],
            "dead": self._stats["dead"],
            "banned": self._stats["banned"],
            "available": self._stats["alive"] - self._stats["banned"],
            "warm": sum(1 for p in self._fast_pool if p.sample_passed),
            "fast_pool": len(self._fast_pool),
            "with_cookies": with_cookies,
            "total_requests_handled": sum(p.requests_handled for p in self._proxies),
            "avg_response_time_s": round(avg_speed, 3),
        }
        self._stats_cache = result
        self._stats_cache_ts = now
        return result


# ─── Singleton ────────────────────────────────────────────────────────────────

_pool: Optional[AsyncProxyPool] = None


def get_proxy_pool() -> AsyncProxyPool:
    global _pool
    if _pool is None:
        _pool = AsyncProxyPool()
    return _pool