agento3/ao3_fetcher.py

"""
异步 AO3 内容抓取器 v5 — 单次代理请求 + CF 挑战检测

v5 vs v4:
- 去掉内部重试循环（重试逻辑迁移到 app.py 的 P2C 选择中）
- 简化 API：直接接受 ProxySession 对象
- 保留 CF 挑战检测 + 状态机反馈
"""
import asyncio
import logging
import time
from typing import Optional

from proxy_pool import ProxySession

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

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

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_',
    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."""
    if status not in (403, 503, 429):
        return False
    server = headers.get("Server", headers.get("server", ""))
    if "cloudflare" not in server.lower():
        for marker in CF_CHALLENGE_MARKERS:
            if marker in body:
                return True
        return False
    for marker in CF_CHALLENGE_MARKERS:
        if marker in body:
            return True
    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/120.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/120.0.0.0 Safari/537.36"
    ),
}

# Timeout config
FAST_REQUEST_TIMEOUT = 8
NORMAL_REQUEST_TIMEOUT = 15

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 _merge_cookies(proxy_cookies: str, user_cookies: str) -> str:
    if not proxy_cookies and not user_cookies:
        return ""
    if not proxy_cookies:
        return user_cookies
    if not user_cookies:
        return proxy_cookies
    return f"{proxy_cookies}; {user_cookies}"


# ─── Main fetch function ──────────────────────────────────────────────────

async def fetch_url(
    url: str,
    proxy: Optional[ProxySession] = None,
    method: str = "GET",
    headers: Optional[dict] = None,
    body: Optional[bytes] = None,
    cookies: Optional[dict] = None,
    is_api: bool = False,
) -> dict:
    """
    Single async fetch through a specific proxy. No internal retry loop.

    Args:
        url: Target AO3 URL
        proxy: ProxySession to use (must already be selected by app.py)
        method: HTTP method
        headers: Request headers
        body: Request body (POST/PUT/PATCH)
        cookies: User cookies dict
        is_api: Use API headers instead of Chrome headers

    Returns:
        dict with success/status/headers/body/cookies/elapsed/is_challenge etc.
    """
    if proxy is None:
        return {
            "status": 0, "headers": {}, "body": b"", "cookies": {},
            "headers_raw": [],
            "success": False, "error": "No proxy provided",
            "elapsed": 0, "proxy_host": None,
            "is_challenge": False, "challenge_body": None, "challenge_proxy": None,
        }

    host_port = proxy.host_port
    start_time = time.time()

    base_headers = API_HEADERS.copy() if is_api else CHROME_HEADERS.copy()
    if headers:
        for h in ["Cookie", "Content-Type", "X-Requested-With",
                   "Accept", "X-CSRF-Token", "Authorization"]:
            if h in headers:
                base_headers[h] = headers[h]
        # Rewrite Origin and Referer: mirror domain → AO3 domain
        # AO3 CSRF checks Origin against archiveofourown.org
        if "Origin" in headers:
            origin = headers["Origin"]
            origin = origin.replace("agento3.miscs.dev", "archiveofourown.org")
            base_headers["Origin"] = origin
        if "Referer" in headers:
            referer = headers["Referer"]
            referer = referer.replace("agento3.miscs.dev", "archiveofourown.org")
            base_headers["Referer"] = referer

    # Filter out OUR cookies — proxy is transparent, only AO3 cookies should reach AO3
    FORWARD_BLOCKED = {"ao3_sessid_proxy", "__cf_bm", "_cfuvid", "_cf_token"}

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

    is_fast = _is_fast_path(url) or method in ("POST", "PUT", "PATCH")
    request_timeout = FAST_REQUEST_TIMEOUT if is_fast else NORMAL_REQUEST_TIMEOUT

    try:
        session = await proxy.get_session()
        request_headers = base_headers.copy()

        # Inject proxy-level cookies (cf_clearance)
        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
        request_cookie_preview = cookie_str[:500] + "..." if len(cookie_str) > 500 else cookie_str
        logger.info(f"[FETCH] {method} {url[:80]} full_cookie={request_cookie_preview}")
        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 = time.time() - start_time
        status = resp.status_code
        resp_body = resp.content
        resp_headers = dict(resp.headers)  # For Content-Type/Location lookups
        resp_headers_raw = list(resp.headers.multi_items())  # Preserves ALL Set-Cookie values

        # Debug: log response Set-Cookie count
        set_cookie_count = sum(1 for k, v in resp_headers_raw if k.lower() == "set-cookie")
        resp_status = resp.status_code
        body_preview = resp_body[:200].decode('utf-8', errors='replace')
        logger.info(f"[FETCH_RESP] {method} {url[:60]} status={resp_status} set_cookies={set_cookie_count} body={body_preview}")

        resp_cookies = {}
        if hasattr(resp, "cookies"):
            for k, v in resp.cookies.items():
                resp_cookies[k] = v

        # CF challenge detection
        if is_cf_challenge(status, resp_body, resp_headers):
            logger.warning(f"[CF_CHALLENGE] {host_port} -> {url[:60]}: {status}")
            return {
                "status": status, "headers": resp_headers, "body": resp_body,
                "headers_raw": resp_headers_raw,
                "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
        if 200 <= status < 500:
            sc_from_dict = resp_headers.get("Set-Cookie", resp_headers.get("set-cookie", "NONE"))
            logger.info(f"[SAVE_COOKIES] proxy={host_port} set_cookie_dict={sc_from_dict[:120]}...")
            proxy.save_cookies(resp_headers)
            return {
                "status": status, "headers": resp_headers, "body": resp_body,
                "headers_raw": resp_headers_raw,
                "cookies": resp_cookies,
                "success": True,
                "elapsed": elapsed, "proxy_host": host_port,
                "is_challenge": False, "challenge_body": None, "challenge_proxy": None,
            }

        # Error status
        logger.warning(f"{host_port} -> {url[:60]}: {status}")
        return {
            "status": status, "headers": resp_headers, "body": resp_body,
            "headers_raw": resp_headers_raw,
            "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:
        elapsed = time.time() - start_time
        logger.warning(f"TIMEOUT: {host_port} -> {url[:60]}")
        return {
            "status": 0, "headers": {}, "body": b"", "cookies": {},
            "headers_raw": [],
            "success": False, "error": "TIMEOUT",
            "elapsed": elapsed, "proxy_host": host_port,
            "is_challenge": False, "challenge_body": None, "challenge_proxy": None,
        }

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