agento3/proxy_pool.py

"""
异步代理池 v5 — P2C + 状态机 + 全量 5000 代理 + 并发限流

对标 go3 的代理管理架构，适配 Python/curl_cffi：
- 全量加载 proxies.txt（不再按端口过滤）
- P2C (Power of Two Choices) 选择算法
- 状态机: healthy → unstable → blocked → probing
- 并发限制 400（Webshare 上限）
- 被动探测 + 指数退避
"""
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")

# ─── Constants ───────────────────────────────────────────────────────────────

PROXY_FILE = "/home/ubuntu/proxy.txt"
WORKING_PROXIES_FILE = "/dev/shm/working_proxies.txt"

# Concurrency limit: Webshare allows ~500 concurrent, cap at 400 for safety
MAX_CONCURRENT_REQUESTS = 400

# Probe config (like go3: 10s interval, 10 concurrency, 2 successes to promote)
PROBE_INTERVAL = 10
PROBE_CONCURRENCY = 10
PROBE_SUCCESS_THRESHOLD = 2
PROBE_TIMEOUT = 5

# Cooldown config
BLOCKED_COOLDOWN_BASE = 120    # First block: 2 min
BLOCKED_COOLDOWN_MAX = 1800    # Max: 30 min
UNSTABLE_COOLDOWN = 15         # Transient failure: 15s

# Outage attack detection
OUTAGE_THRESHOLD = 0.9  # 90% proxies unavailable → attack page

# Max retries for a single request
MAX_RETRIES = 5

# Ban-worthy status codes
BAN_STATUS_CODES = {403, 525}

# TLS fingerprints (proven against AO3 CF)
BROWSER_IMPS = ["safari15_5", "safari17_0", "chrome123", "chrome124"]

# States
STATE_HEALTHY = "healthy"
STATE_UNSTABLE = "unstable"
STATE_BLOCKED = "blocked"
STATE_PROBING = "probing"


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


def load_proxies_from_file() -> list[str]:
    """Load ALL proxies from file. No port filtering in v5."""
    # Prefer cached working proxies if available and fresh
    # But only use as a hint — still load the full list
    if os.path.exists(WORKING_PROXIES_FILE) and os.path.getsize(WORKING_PROXIES_FILE) > 0:
        with open(WORKING_PROXIES_FILE) as f:
            cached = [l.strip() for l in f if l.strip() and ":" in l]
        if len(cached) >= 100:
            logger.info(f"Found {len(cached)} cached working proxies — using as seed")

    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)

    if not proxies:
        raise RuntimeError(f"No proxies found in {PROXY_FILE}")

    logger.info(f"Loaded {len(proxies)} proxies from {PROXY_FILE}")
    return proxies


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

    __slots__ = (
        "idx", "host_port", "host", "port_str", "port",
        "session", "impersonate",
        # State machine
        "state",  # healthy / unstable / blocked / probing
        # Stats
        "successes", "failures", "consecutive_net_failures",
        "consecutive_blocked",
        "last_used", "last_success_at", "last_failure_at",
        "avg_response_time", "requests_handled",
        # Cooldown / block
        "cooldown_until", "next_probe_at",
        "probe_successes",
        # v5: inflight counter
        "inflight",
        # Cookie jar
        "_cookies", "_cookie_expires", "_lock",
    )

    def __init__(self, idx: int, host_port: str):
        self.idx = idx
        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.state = STATE_HEALTHY
        self.successes = 0
        self.failures = 0
        self.consecutive_net_failures = 0
        self.consecutive_blocked = 0
        self.last_used = 0.0
        self.last_success_at = 0.0
        self.last_failure_at = 0.0
        self.avg_response_time = 1.0
        self.requests_handled = 0
        self.cooldown_until = 0.0
        self.next_probe_at = 0.0
        self.probe_successes = 0
        self.inflight = 0

        self._cookies: dict[str, str] = {}
        self._cookie_expires: dict[str, float] = {}
        self._lock = Lock()

    @property
    def is_available(self) -> bool:
        """Proxy is selectable for user requests."""
        now = time.time()
        if now < self.cooldown_until:
            return False
        if self.state == STATE_HEALTHY:
            return True
        if self.state == STATE_UNSTABLE and now >= self.cooldown_until:
            return True
        if self.state == STATE_PROBING:
            return False  # Don't route user traffic to probing proxies
        return False  # STATE_BLOCKED

    @property
    def is_user_selectable(self) -> bool:
        """Proxy can be chosen for sticky/user requests. Like go3's isUserSelectable."""
        return self.is_available

    def _purge_expired(self, now: float):
        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)

    def save_cookies(self, headers: dict) -> int:
        """Extract ONLY cf_clearance from response headers. Returns count saved.

        The proxy must be TRANSPARENT to AO3 — do NOT save _otwarchive_session,
        user_credentials, or any other AO3 cookies. Only cf_clearance (Cloudflare
        bypass) is our concern. Everything else passes through to the browser.
        """
        set_cookie = headers.get("Set-Cookie", headers.get("set-cookie", ""))
        if not set_cookie:
            return 0
        now = time.time()
        with self._lock:
            parts = set_cookie.split(";")
            if parts:
                first_part = parts[0].strip()
                if "=" in first_part:
                    name, _, value = first_part.partition("=")
                    name = name.strip()
                    value = value.strip()
                    # ONLY save cf_clearance — the only cookie the proxy needs
                    if name == "cf_clearance":
                        self._cookies[name] = value
                        logger.info(f"[CF_SAVE] {self.host_port}: saved cf_clearance")
                        expires = _parse_cookie_expires(set_cookie)
                        if expires > 0:
                            self._cookie_expires[name] = expires
                        self._purge_expired(now)
                        return 1
                    else:
                        logger.info(f"[CF_SKIP] {self.host_port}: skipping {name} (not cf_clearance)")
            self._purge_expired(now)
        return 0

    def get_cookie_header(self) -> str:
        now = time.time()
        with self._lock:
            self._purge_expired(now)
            if not self._cookies:
                return ""
            header = "; ".join(f"{k}={v}" for k, v in self._cookies.items())
            logger.info(f"[PROXY_COOKIES] {self.host_port}: {list(self._cookies.keys())[:10]}")
            return header

    def has_cookies(self) -> bool:
        return bool(self._cookies)

    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}",
            }
            self.session.headers.update({
                "User-Agent": (
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                    "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.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

    # ─── State machine (go3-compatible) ──────────────────────────────────

    def mark_success(self, response_time: float):
        """Request succeeded. Promote to healthy."""
        self.last_used = time.time()
        self.last_success_at = time.time()
        self.successes += 1
        self.requests_handled += 1
        self.consecutive_net_failures = 0
        self.consecutive_blocked = 0
        self.avg_response_time = self.avg_response_time * 0.7 + response_time * 0.3
        self.cooldown_until = 0.0
        self.state = STATE_HEALTHY

    def mark_failure(self, is_network: bool = True):
        """Connection error / timeout. Exponential backoff."""
        now = time.time()
        self.last_used = now
        self.last_failure_at = now
        self.failures += 1
        self.requests_handled += 1

        if is_network:
            self.consecutive_net_failures += 1
            self.consecutive_blocked = 0

        if self.consecutive_net_failures >= 3:
            self.state = STATE_BLOCKED
            # Exponential backoff: base * 2^(failures-1), capped at max
            backoff = min(
                BLOCKED_COOLDOWN_BASE * (2 ** (self.consecutive_net_failures - 3)),
                BLOCKED_COOLDOWN_MAX
            )
            self.cooldown_until = now + backoff
            self.next_probe_at = now + backoff
            self.probe_successes = 0
        elif self.consecutive_net_failures >= 1:
            self.state = STATE_UNSTABLE
            self.cooldown_until = now + UNSTABLE_COOLDOWN
        # else stay healthy / unchanged

    def mark_blocked(self):
        """403/525 — CF blocked. Move to blocked state."""
        now = time.time()
        self.last_used = now
        self.last_failure_at = now
        self.failures += 1
        self.requests_handled += 1
        self.consecutive_blocked += 1
        self.consecutive_net_failures = 0

        self.state = STATE_BLOCKED
        backoff = min(
            BLOCKED_COOLDOWN_BASE * (2 ** (self.consecutive_blocked - 1)),
            BLOCKED_COOLDOWN_MAX
        )
        self.cooldown_until = now + backoff
        self.next_probe_at = now + backoff
        self.probe_successes = 0

    def mark_challenged(self):
        """CF challenge detected — proxy is alive but needs cookie. Keep healthy, don't penalize."""
        self.last_used = time.time()
        self.requests_handled += 1
        # Don't change state — challenge is not proxy's fault

    def mark_probe_result(self, passed: bool, response_time: float = 0):
        """Periodic probe result. If probing and enough successes, promote to healthy."""
        now = time.time()
        if passed:
            self.probe_successes += 1
            if self.probe_successes >= PROBE_SUCCESS_THRESHOLD:
                self.state = STATE_HEALTHY
                self.cooldown_until = 0.0
                self.consecutive_net_failures = 0
                self.consecutive_blocked = 0
                self.avg_response_time = self.avg_response_time * 0.5 + response_time * 0.5
        else:
            self.probe_successes = 0
            # Extend cooldown
            self.cooldown_until = now + max(self.cooldown_until - now, BLOCKED_COOLDOWN_BASE)
        self.next_probe_at = 0.0

    def to_stats_dict(self) -> dict:
        return {
            "address": self.host_port,
            "state": self.state,
            "successes": self.successes,
            "failures": self.failures,
            "consecutive_net_failures": self.consecutive_net_failures,
            "consecutive_blocked": self.consecutive_blocked,
            "last_used": self.last_used,
            "last_success_at": self.last_success_at,
            "last_failure_at": self.last_failure_at,
            "avg_response_time": round(self.avg_response_time, 3),
            "requests_handled": self.requests_handled,
            "cooldown_until": self.cooldown_until,
            "inflight": self.inflight,
            "has_cookies": bool(self._cookies),
            "impersonate": self.impersonate,
        }

    def __repr__(self):
        return f"PS({self.host_port}, {self.state}, {self.avg_response_time:.2f}s, inflight={self.inflight})"


class AsyncProxyPool:
    """Proxy pool v5 — P2C selection + state machine + concurrent limiter."""

    def __init__(self):
        self._proxies: list[ProxySession] = []
        self._idx_map: dict[int, ProxySession] = {}
        self._addr_map: dict[str, ProxySession] = {}

        # Concurrent request limiter (capped at 400)
        self._concurrency_limiter = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)

        # Stats
        self._stats_lock = Lock()

        # Probe task
        self._probe_task: Optional[asyncio.Task] = None
        self._probe_running = False

        # Load proxies
        self._load_proxies()
        self._start_probing()

        logger.info(
            f"ProxyPool v5 ready: {len(self._proxies)} proxies, "
            f"concurrency={MAX_CONCURRENT_REQUESTS}, selection=p2c"
        )

    def _load_proxies(self):
        proxy_list = load_proxies_from_file()
        self._proxies = []
        self._idx_map = {}
        self._addr_map = {}
        for i, hp in enumerate(proxy_list):
            ps = ProxySession(i, hp)
            self._proxies.append(ps)
            self._idx_map[i] = ps
            self._addr_map[hp] = ps
        logger.info(f"Loaded {len(self._proxies)} proxies")

    def _start_probing(self):
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
        self._probe_task = asyncio.create_task(self._probe_loop())

    # ─── Proxy selection (P2C — Power of Two Choices) ───────────────────

    def p2c_select(self, exclude_idxs: Optional[set[int]] = None) -> Optional[ProxySession]:
        """
        Power of Two Choices: pick 2 random candidates, return the one with fewer inflight.
        Falls back to random if P2C produces ties.
        """
        exclude = exclude_idxs or set()
        now = time.time()

        # Build candidate list (selectable, not excluded)
        candidates = []
        for p in self._proxies:
            if p.idx in exclude:
                continue
            if not p.is_user_selectable:
                continue
            if p.inflight >= 10:  # Don't overload a single proxy
                continue
            candidates.append(p)

        if not candidates:
            # Fallback: anything with cooldown expired, even if probing
            for p in self._proxies:
                if p.idx in exclude:
                    continue
                if now >= p.cooldown_until and p.inflight < 10:
                    candidates.append(p)

        if not candidates:
            # Desperate: any proxy
            for p in self._proxies:
                if p.idx in exclude:
                    continue
                if p.inflight < 10:
                    candidates.append(p)

        if not candidates:
            return None

        if len(candidates) == 1:
            return candidates[0]

        # P2C: pick 2, choose the one with lower inflight
        left = random.choice(candidates)
        right = random.choice(candidates)
        # Avoid same proxy
        for _ in range(5):
            if right.idx != left.idx:
                break
            right = random.choice(candidates)

        if left.inflight <= right.inflight:
            return left
        return right

    def get_by_idx(self, idx: int) -> Optional[ProxySession]:
        return self._idx_map.get(idx)

    def get_by_addr(self, addr: str) -> Optional[ProxySession]:
        return self._addr_map.get(addr)

    # ─── Sticky proxy ───────────────────────────────────────────────────

    def sticky_select(self, sticky_idx: int, exclude_idxs: set[int]) -> Optional[ProxySession]:
        """Try to use the sticky proxy if it's still selectable."""
        p = self.get_by_idx(sticky_idx)
        if p and p.is_user_selectable and p.idx not in exclude_idxs and p.inflight < 10:
            return p
        return None

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

    def get_stats(self) -> dict:
        total = len(self._proxies)
        healthy = 0
        unstable = 0
        blocked = 0
        probing = 0
        total_inflight = 0
        with_cookies = 0
        total_requests = 0
        response_times = []

        for p in self._proxies:
            if p.state == STATE_HEALTHY:
                healthy += 1
            elif p.state == STATE_UNSTABLE:
                unstable += 1
            elif p.state == STATE_BLOCKED:
                blocked += 1
            elif p.state == STATE_PROBING:
                probing += 1
            total_inflight += p.inflight
            if p.has_cookies():
                with_cookies += 1
            total_requests += p.requests_handled
            if p.requests_handled > 0:
                response_times.append(p.avg_response_time)

        avg_rt = sum(response_times) / max(len(response_times), 1)

        return {
            "total": total,
            "healthy": healthy,
            "unstable": unstable,
            "blocked": blocked,
            "probing": probing,
            "available": healthy,  # Only healthy are truly available
            "total_inflight": total_inflight,
            "with_cookies": with_cookies,
            "total_requests": total_requests,
            "avg_response_time_s": round(avg_rt, 3),
            "pool_health": round(healthy / max(total, 1), 3),
            "pool_unavailable_ratio": round((blocked + probing) / max(total, 1), 3),
        }

    def get_proxy_stats_list(self) -> list[dict]:
        """Detailed per-proxy stats for the monitor."""
        return [p.to_stats_dict() for p in self._proxies]

    def pool_available_ratio(self) -> float:
        """Fraction of proxies currently selectable."""
        available = sum(1 for p in self._proxies if p.is_user_selectable)
        return available / max(len(self._proxies), 1)

    def pool_unavailable_ratio(self) -> float:
        return 1.0 - self.pool_available_ratio()

    # ─── Concurrency limiter ────────────────────────────────────────────

    @property
    def concurrency_limiter(self) -> asyncio.Semaphore:
        return self._concurrency_limiter

    # ─── Probe loop ─────────────────────────────────────────────────────

    async def _probe_loop(self):
        """Periodic probe to recover blocked proxies. Like go3's probe goroutine."""
        await asyncio.sleep(5)  # Wait for initial startup
        logger.info("Probe loop started")
        while True:
            try:
                await asyncio.sleep(PROBE_INTERVAL)
                # Find proxies that need probing: blocked/probing with cooldown expired
                now = time.time()
                to_probe = []
                for p in self._proxies:
                    if p.state in (STATE_BLOCKED, STATE_PROBING):
                        if now >= p.cooldown_until and now >= p.next_probe_at:
                            to_probe.append(p)

                if not to_probe:
                    continue

                # Probe in batches of PROBE_CONCURRENCY
                random.shuffle(to_probe)
                for i in range(0, len(to_probe), PROBE_CONCURRENCY):
                    batch = to_probe[i:i + PROBE_CONCURRENCY]
                    tasks = [self._probe_single(p) for p in batch]
                    results = await asyncio.gather(*tasks, return_exceptions=True)
                    for p, r in zip(batch, results):
                        if isinstance(r, Exception):
                            p.mark_probe_result(False)
                        else:
                            p.mark_probe_result(r[0], r[1])

                # Log recovery stats
                recovered = sum(1 for p in to_probe if p.state == STATE_HEALTHY)
                if recovered:
                    logger.info(f"Probe: {len(to_probe)} checked, {recovered} recovered")
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Probe error: {e}")

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

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

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

    @property
    def proxies(self) -> list[ProxySession]:
        return self._proxies

    def __len__(self):
        return len(self._proxies)


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

_pool: Optional[AsyncProxyPool] = None


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