agento3/url_rewriter.py

"""
URL 重写器
- HTML 内容中所有 archiveofourown.org 替换为 agento3.miscs.dev
- CSS/JS 中的 URL 替换
- 响应头中的 Location/Set-Cookie 重写
"""

import re
from typing import Optional

TARGET_DOMAIN = "archiveofourown.org"
MIRROR_DOMAIN = "agento3.miscs.dev"

# Pre-compiled regex for efficiency
DOMAIN_RE = re.compile(rb"archiveofourown\.org", re.IGNORECASE)
# Match URLs in HTML/CSS/JS
URL_RE = re.compile(
    rb'(https?://)(www\.)?archiveofourown\.org',
    re.IGNORECASE
)

# Content types that need body rewriting
REWRITABLE_CONTENT_TYPES = {
    "text/html",
    "text/plain",
    "text/css",
    "application/javascript",
    "application/x-javascript",
    "text/javascript",
    "application/json",
    "application/xml",
    "text/xml",
    "application/atom+xml",
    "application/rss+xml",
}

# CSS-specific patterns
CSS_URL_RE = re.compile(rb'url\([\'"]?(https?://[^\)\'"]*archiveofourown\.org[^\)\'"]*)[\'"]?\)', re.IGNORECASE)

# JS-specific patterns (strings containing the domain)
JS_DOMAIN_RE = re.compile(rb'["\'](https?://[^"\']*archiveofourown\.org[^"\']*)["\']', re.IGNORECASE)


def needs_rewrite(content_type: str) -> bool:
    """Check if this content type needs body rewriting."""
    if not content_type:
        return False
    ct = content_type.split(";")[0].strip().lower()
    return ct in REWRITABLE_CONTENT_TYPES


def rewrite_body(body: bytes, content_type: Optional[str] = None) -> bytes:
    """
    Rewrite AO3 URLs in body content.
    If content_type is provided, only rewrite if it's a rewritable type.
    """
    if content_type and not needs_rewrite(content_type):
        return body

    # Simple domain replacement
    return DOMAIN_RE.sub(MIRROR_DOMAIN.encode(), body)


def rewrite_response_headers(headers: dict) -> dict:
    """Rewrite AO3 domains in response headers."""
    new_headers = {}
    for key, value in headers.items():
        key_lower = key.lower()

        if key_lower == "location":
            # Redirect targets
            value = value.replace(TARGET_DOMAIN, MIRROR_DOMAIN)
            # Also rewrite protocol if needed
            value = value.replace("http://", "https://")

        elif key_lower == "set-cookie":
            # Cookie domain rewrite
            value = value.replace(f"domain={TARGET_DOMAIN}", f"domain={MIRROR_DOMAIN}")
            value = value.replace(f"Domain={TARGET_DOMAIN}", f"Domain={MIRROR_DOMAIN}")

        new_headers[key] = value

    return new_headers


def rewrite_response_headers_raw(headers_raw: list) -> list:
    """Rewrite AO3 domains in raw response headers (preserves multi-value Set-Cookie).

    Unlike rewrite_response_headers which works on a dict (collapsing multi-value
    Set-Cookie), this operates on a list of (key, value) tuples so ALL Set-Cookie
    values are preserved and individually rewritten.
    """
    result = []
    for key, value in headers_raw:
        key_lower = key.lower()

        if key_lower == "location":
            value = value.replace(TARGET_DOMAIN, MIRROR_DOMAIN)
            value = value.replace("http://", "https://")

        elif key_lower == "set-cookie":
            value = value.replace(f"domain={TARGET_DOMAIN}", f"domain={MIRROR_DOMAIN}")
            value = value.replace(f"Domain={TARGET_DOMAIN}", f"Domain={MIRROR_DOMAIN}")

        result.append((key, value))

    return result


def rewrite_redirect_url(url: str) -> str:
    """Rewrite AO3 URLs in redirect targets."""
    return url.replace(TARGET_DOMAIN, MIRROR_DOMAIN).replace("http://", "https://")