agento3/ao3_fetcher.py

"""
异步 AO3 内容抓取器 v4 — Cookie 感知 + CF 挑战检测 + 智能重试

v4 vs v3:
- CF 挑战检测：不再把 403/503 一律当失败
- Cookie 注入：请求时自动携带 proxy 级别的 cookies（cf_clearance 等）
- Cookie 保存：成功请求后自动提取 Set-Cookie 并保存到 proxy cookie jar
- 智能重试：挑战时优先用带 cookie 的代理，无 cookie 则换代理
- 挑战页面透传：所有重试都遇到挑战时，返回挑战 HTML 让用户浏览器求解
"""
import asyncio
import logging
import time
from typing import Optional

from proxy_pool import get_proxy_pool, ProxySession

logger = logging.getLogger("ao3-fetcher")

# ─── CF Challenge Detection ───────────────────────────────────────────────

# Markers in response body that indicate a Cloudflare challenge page
CF_CHALLENGE_MARKERS = [
    b'/cdn-cgi/challenge-platform',
    b'cf-challenge-running',
    b'cf-browser-verification',
    b'window._cf_chl_opt',
    b'challenge-platform',
    b'cf-turnstile',
    b'cf_chl_',
    # Sometimes CF just returns "Just a moment..." without JS markers
    b'Checking your browser',
    b'Just a moment...',
]


def is_cf_challenge(status: int, body: bytes, headers: dict) -> bool:
    """Detect if response is a Cloudflare challenge page (not a real error)."""
    if status not in (403, 503, 429):
        return False
    # Quick check: CF always sets Server header on challenge pages
    server = headers.get("Server", headers.get("server", ""))
    if "cloudflare" not in server.lower():
        # Check body for challenge markers
        for marker in CF_CHALLENGE_MARKERS:
            if marker in body:
                return True
        return False
    # Server: cloudflare + non-200 status = likely challenge
    for marker in CF_CHALLENGE_MARKERS:
        if marker in body:
            return True
    # If Server is cloudflare and status is 403/503, it's almost certainly a challenge
    if status in (403, 503):
        return True
    return False


# ─── Headers ──────────────────────────────────────────────────────────────

CHROME_HEADERS = {
    "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,zh-CN;q=0.8,zh;q=0.7",
    "Sec-Ch-Ua": '"Not A(Brand";v="99", "Google Chrome";v="120", "Chromium";v="120"',
    "Sec-Ch-Ua-Mobile": "?0",
    "Sec-Ch-Ua-Platform": '"Windows"',
    "Sec-Fetch-Dest": "document",
    "Sec-Fetch-Mode": "navigate",
    "Sec-Fetch-Site": "none",
    "Sec-Fetch-User": "?1",
    "Upgrade-Insecure-Requests": "1",
    "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",
    "DNT": "1",
}

API_HEADERS = {
    "Accept": "*/*",
    "Accept-Language": "en-US,en;q=0.9",
    "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",
}

# ─── Timeout config ───────────────────────────────────────────────────────

FAST_REQUEST_TIMEOUT = 8   # login/register/signup
NORMAL_REQUEST_TIMEOUT = 15
CONNECT_TIMEOUT = 5
MAX_RETRIES = 2
FAST_MAX_RETRIES = 1

FAST_PATHS = {
    "/users/login",
    "/users/sign_up",
    "/users/new",
    "/invitation_requests",
    "/token_dispenser.json",
    "/user_sessions",
}


def _is_fast_path(url: str) -> bool:
    for fp in FAST_PATHS:
        if fp in url:
            return True
    if "/users/" in url or "/user_sessions" in url:
        return True
    return False


def _now() -> float:
    try:
        return asyncio.get_running_loop().time()
    except RuntimeError:
        return time.time()


def _merge_cookies(proxy_cookies: str, user_cookies: str) -> str:
    """Merge proxy-level cookies (cf_clearance) with user cookies. User cookies take precedence."""
    if not proxy_cookies and not user_cookies:
        return ""
    if not proxy_cookies:
        return user_cookies
    if not user_cookies:
        return proxy_cookies
    # User cookies take priority (they contain session auth)
    # But put proxy cookies first so user cookies can override
    return f"{proxy_cookies}; {user_cookies}"


async def fetch_url(
    url: str,
    method: str = "GET",
    headers: Optional[dict] = None,
    body: Optional[bytes] = None,
    cookies: Optional[dict] = None,
    is_api: bool = False,
    preferred_proxy: Optional[str] = None,  # v4: proxy host:port for challenge affinity
) -> dict:
    """
    Async fetch from AO3 through proxy pool with cookie-aware session reuse.

    Priority-based:
    - Login/register paths → fast pool + 8s timeout + 1 retry
    - Normal paths → cookie-preferring proxies + 15s timeout + 2 retries

    v4 improvements:
    - CF challenge detection: don't treat challenge as proxy failure
    - Cookie injection: auto-send saved cf_clearance per proxy
    - Cookie saving: auto-extract Set-Cookie on success
    - Challenge body returned for user-browser solving

    Returns: {
        "status": 200,
        "headers": {...},
        "body": b"...",
        "cookies": {...},
        "success": True/False,
        "error": "...",
        "elapsed": 0.5,
        "proxy_host": "p.webshare.io:10296",
        # v4:
        "is_challenge": False,
        "challenge_body": b"" | None,  # CF challenge HTML for user-browser solving
        "challenge_proxy": "" | None,  # which proxy got the challenge
    }
    """
    pool = get_proxy_pool()
    base_headers = API_HEADERS.copy() if is_api else CHROME_HEADERS.copy()

    if headers:
        for h in ["Cookie", "Referer", "Content-Type", "X-Requested-With", "Accept",
                   "Origin", "X-CSRF-Token", "Authorization"]:
            if h in headers:
                base_headers[h] = headers[h]

    user_cookie_str = ""
    if cookies:
        user_cookie_str = "; ".join(f"{k}={v}" for k, v in cookies.items())

    # Determine priority level
    is_fast = _is_fast_path(url) or method in ("POST", "PUT", "PATCH")
    request_timeout = FAST_REQUEST_TIMEOUT if is_fast else NORMAL_REQUEST_TIMEOUT
    max_retries = FAST_MAX_RETRIES if is_fast else MAX_RETRIES

    last_error = None
    last_challenge_body = None
    last_challenge_proxy = None
    seen_proxies = set()  # Don't retry with same proxy

    # v4: if preferred_proxy is set (from challenge affinity), use it first
    if preferred_proxy:
        try:
            target_proxy = None
            for p in pool._proxies:
                if p.host_port == preferred_proxy:
                    target_proxy = p
                    break
            if target_proxy and target_proxy.is_available:
                result = await _try_proxy(
                    target_proxy, url, method, body, base_headers, user_cookie_str,
                    request_timeout, is_fast
                )
                if result["success"]:
                    # preferred proxy worked — save cookies
                    target_proxy.save_cookies(result.get("headers", {}))
                    return result
                if result.get("is_challenge"):
                    last_challenge_body = result.get("challenge_body")
                    last_challenge_proxy = target_proxy.host_port
                seen_proxies.add(target_proxy.host_port)
        except Exception:
            pass

    for attempt in range(max_retries):
        # v4: prefer proxies with cookies, fallback to any available
        if attempt == 0:
            # First attempt: use proxy with cookies if available
            proxy = pool.get_proxy_with_cookies()
            if not proxy:
                proxy = pool.get_fast_proxy() if is_fast else pool.get_proxy()
        elif attempt == 1 and last_challenge_body:
            # Second attempt after challenge: try another proxy with cookies
            proxy = pool.get_proxy_with_cookies()
            # Make sure it's not the same one
            if proxy and proxy.host_port in seen_proxies:
                proxy = pool.get_fast_proxy() if is_fast else pool.get_proxy()
            if not proxy:
                proxy = pool.get_fast_proxy() if is_fast else pool.get_proxy()
        else:
            proxy = pool.get_fast_proxy() if is_fast else pool.get_proxy()

        if not proxy or proxy.host_port in seen_proxies:
            # Find an unseen proxy
            for _ in range(5):
                p = pool.get_proxy()
                if p and p.host_port not in seen_proxies:
                    proxy = p
                    break
            if not proxy or proxy.host_port in seen_proxies:
                continue

        seen_proxies.add(proxy.host_port)

        result = await _try_proxy(
            proxy, url, method, body, base_headers, user_cookie_str,
            request_timeout, is_fast
        )

        if result["success"]:
            # Save response cookies to proxy cookie jar for future requests
            proxy.save_cookies(result.get("headers", {}))
            return result

        if result.get("is_challenge"):
            last_challenge_body = result.get("challenge_body")
            last_challenge_proxy = proxy.host_port
            last_error = result.get("error", "CF_CHALLENGE")
            # Don't break — try another proxy
            continue

        # Real failure (connection error, timeout)
        last_error = result.get("error", "UNKNOWN")

    # All retries exhausted
    return {
        "status": 0, "headers": {}, "body": b"", "cookies": {},
        "success": False,
        "error": f"All retries failed: {last_error}",
        "elapsed": 0, "proxy_host": None,
        "is_challenge": last_challenge_body is not None,
        "challenge_body": last_challenge_body,
        "challenge_proxy": last_challenge_proxy,
    }


async def _try_proxy(
    proxy: ProxySession,
    url: str,
    method: str,
    body: Optional[bytes],
    base_headers: dict,
    user_cookie_str: str,
    request_timeout: int,
    is_fast: bool,
) -> dict:
    """Try a single request through one proxy. Returns result dict."""
    host_port = f"{proxy.host}:{proxy.port}"
    start_time = _now()

    try:
        session = await proxy.get_session()

        # v4: Build per-request headers with proxy cookies
        # IMPORTANT: use headers= parameter (never session.headers.update() — race condition!)
        request_headers = base_headers.copy()

        # Inject proxy-level cookies (cf_clearance, etc.)
        proxy_cookie_str = proxy.get_cookie_header()
        cookie_str = _merge_cookies(proxy_cookie_str, user_cookie_str)
        if cookie_str:
            request_headers["Cookie"] = cookie_str

        # Execute request
        if method == "GET":
            resp = await session.get(url, timeout=request_timeout, headers=request_headers)
        elif method == "POST":
            resp = await session.post(url, data=body, timeout=request_timeout, headers=request_headers)
        elif method == "HEAD":
            resp = await session.head(url, timeout=request_timeout, headers=request_headers)
        else:
            resp = await session.request(method, url, data=body, timeout=request_timeout,
                                         headers=request_headers)

        elapsed = _now() - start_time
        status = resp.status_code
        resp_body = resp.content
        resp_headers = dict(resp.headers)
        resp_cookies = {}
        if hasattr(resp, "cookies"):
            for k, v in resp.cookies.items():
                resp_cookies[k] = v

        # v4: Check for CF challenge
        if is_cf_challenge(status, resp_body, resp_headers):
            proxy.mark_challenged()
            logger.warning(f"[CF_CHALLENGE] {host_port} -> {url[:60]}: {status} ({elapsed:.1f}s)")
            return {
                "status": status,
                "headers": resp_headers,
                "body": resp_body,
                "cookies": resp_cookies,
                "success": False,
                "error": f"CF_CHALLENGE_{status}",
                "elapsed": elapsed,
                "proxy_host": host_port,
                "is_challenge": True,
                "challenge_body": resp_body,
                "challenge_proxy": host_port,
            }

        # Success — any 2xx-4xx is proxied through
        if 200 <= status < 500:
            proxy.mark_success(elapsed)
            logger.debug(f"{host_port} -> {url[:60]}: {status} ({elapsed:.2f}s)")

            # v4: Save response cookies to proxy
            proxy.save_cookies(resp_headers)

            return {
                "status": status,
                "headers": resp_headers,
                "body": resp_body,
                "cookies": resp_cookies,
                "success": True,
                "elapsed": elapsed,
                "proxy_host": host_port,
                "is_challenge": False,
                "challenge_body": None,
                "challenge_proxy": None,
            }

        # True error status (5xx)
        proxy.mark_failure()
        logger.warning(f"Attempt: {host_port} -> {url[:60]}: {status}")
        return {
            "status": status,
            "headers": resp_headers,
            "body": resp_body,
            "cookies": resp_cookies,
            "success": False,
            "error": f"HTTP_{status}",
            "elapsed": elapsed,
            "proxy_host": host_port,
            "is_challenge": False,
            "challenge_body": None,
            "challenge_proxy": None,
        }

    except asyncio.TimeoutError:
        proxy.mark_failure()
        elapsed = _now() - start_time
        logger.warning(f"TIMEOUT: {host_port} -> {url[:60]} ({elapsed:.1f}s, timeout={request_timeout}s)")
        return {
            "status": 0, "headers": {}, "body": b"", "cookies": {},
            "success": False, "error": "TIMEOUT",
            "elapsed": elapsed, "proxy_host": host_port,
            "is_challenge": False, "challenge_body": None, "challenge_proxy": None,
        }

    except Exception as e:
        proxy.mark_failure()
        elapsed = _now() - start_time
        err_str = str(e)[:120]
        logger.debug(f"ERROR: {host_port} -> {e} ({elapsed:.1f}s)")
        return {
            "status": 0, "headers": {}, "body": b"", "cookies": {},
            "success": False, "error": err_str,
            "elapsed": elapsed, "proxy_host": host_port,
            "is_challenge": False, "challenge_body": None, "challenge_proxy": None,
        }