proxy-god/proxy_chain_manager/ban_tester.py

from __future__ import annotations

import concurrent.futures
import socket
from dataclasses import dataclass
from typing import Iterable

import httpx

# ── Site lists ────────────────────────────────────────────────────────────────

SITES_SOCIAL: tuple[tuple[str, str], ...] = (
    ("Google",      "https://www.google.com/generate_204"),
    ("YouTube",     "https://www.youtube.com/"),
    ("Facebook",    "https://www.facebook.com/"),
    ("Instagram",   "https://www.instagram.com/"),
    ("X / Twitter", "https://x.com/"),
    ("Reddit",      "https://www.reddit.com/"),
    ("TikTok",      "https://www.tiktok.com/"),
    ("Wikipedia",   "https://www.wikipedia.org/"),
    ("DuckDuckGo",  "https://duckduckgo.com/"),
    ("BBC",         "https://www.bbc.com/"),
    ("GitHub",      "https://github.com/"),
)

SITES_SHOPPING: tuple[tuple[str, str], ...] = (
    ("Amazon",      "https://www.amazon.com/"),
    ("eBay",        "https://www.ebay.com/"),
    ("Walmart",     "https://www.walmart.com/"),
    ("Etsy",        "https://www.etsy.com/"),
    ("AliExpress",  "https://www.aliexpress.com/"),
    ("Best Buy",    "https://www.bestbuy.com/"),
    ("Target",      "https://www.target.com/"),
    ("Newegg",      "https://www.newegg.com/"),
)

SITES_CRYPTO: tuple[tuple[str, str], ...] = (
    ("Coinbase",    "https://www.coinbase.com/"),
    ("Binance",     "https://www.binance.com/"),
    ("Kraken",      "https://www.kraken.com/"),
    ("Bybit",       "https://www.bybit.com/"),
    ("OKX",         "https://www.okx.com/"),
    ("KuCoin",      "https://www.kucoin.com/"),
    ("Bitfinex",    "https://www.bitfinex.com/"),
    ("Gemini",      "https://www.gemini.com/"),
)

SITES_DNS: tuple[tuple[str, str], ...] = (
    ("Cloudflare DNS (1.1.1.1)", "https://1.1.1.1/"),
    ("Google DNS (8.8.8.8)",     "https://dns.google/"),
    ("Quad9 DNS (9.9.9.9)",      "https://quad9.net/"),
    ("AdGuard DNS",              "https://adguard-dns.io/"),
    ("NextDNS",                  "https://nextdns.io/"),
    ("Mullvad DNS",              "https://mullvad.net/en/help/dns-over-https-and-dns-over-tls/"),
    ("OpenDNS",                  "https://www.opendns.com/"),
)

# Combined default: all categories
POPULAR_SITES: tuple[tuple[str, str], ...] = (
    SITES_SOCIAL + SITES_SHOPPING + SITES_CRYPTO + SITES_DNS
)

SITE_CATEGORIES: dict[str, tuple[tuple[str, str], ...]] = {
    "Social / General": SITES_SOCIAL,
    "Shopping":         SITES_SHOPPING,
    "Crypto Exchanges": SITES_CRYPTO,
    "DNS Providers":    SITES_DNS,
    "All":              POPULAR_SITES,
}

_BANNED_STATUS_CODES = {401, 403, 407, 418, 429, 451}
_BANNED_TEXT_HINTS = (
    "access denied",
    "forbidden",
    "temporarily blocked",
    "request blocked",
    "unusual traffic",
    "captcha",
    "challenge required",
    "bot detected",
    "suspicious activity",
    "geo-blocked",
    "not available in your region",
    "your ip",
    "ip address",
    "cloudflare ray",
    "ddos protection",
    "attention required",
)


@dataclass(frozen=True)
class BanTestResult:
    site: str
    url: str
    status: str      # ok | banned | error
    code: int | None
    detail: str
    category: str = ""


def _looks_banned_body(body: str) -> bool:
    t = (body or "").lower()
    return any(h in t for h in _BANNED_TEXT_HINTS)


def _test_one(
    proxy_url: str,
    site: str,
    url: str,
    timeout_seconds: float,
    category: str = "",
) -> BanTestResult:
    timeout = httpx.Timeout(timeout_seconds, connect=min(8.0, timeout_seconds))
    try:
        with httpx.Client(
            proxy=proxy_url,
            timeout=timeout,
            verify=False,
            follow_redirects=True,
            headers={
                "User-Agent": (
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                    "AppleWebKit/537.36 (KHTML, like Gecko) "
                    "Chrome/124.0.0.0 Safari/537.36"
                ),
                "Accept": "text/html,application/xhtml+xml,*/*;q=0.8",
                "Accept-Language": "en-US,en;q=0.9",
            },
        ) as c:
            r = c.get(url)
            body = r.text[:2400] if r.text else ""
            if r.status_code in _BANNED_STATUS_CODES or _looks_banned_body(body):
                return BanTestResult(
                    site, url, "banned", r.status_code,
                    f"HTTP {r.status_code}", category,
                )
            if 200 <= r.status_code < 400:
                return BanTestResult(site, url, "ok", r.status_code, f"HTTP {r.status_code}", category)
            return BanTestResult(site, url, "error", r.status_code, f"HTTP {r.status_code}", category)
    except Exception as e:
        short = str(e)[:80]
        return BanTestResult(site, url, "error", None, f"{type(e).__name__}: {short}", category)


def _dns_resolve_via_system(hostname: str, timeout: float = 5.0) -> str | None:
    """Resolve a hostname using system DNS (does not go through proxy — exposes leak)."""
    try:
        socket.setdefaulttimeout(timeout)
        info = socket.getaddrinfo(hostname, None)
        for entry in info:
            addr = entry[4][0]
            if addr:
                return addr
    except Exception:
        pass
    return None


@dataclass(frozen=True)
class DnsLeakResult:
    resolver: str
    ip_via_proxy: str | None
    ip_direct: str | None
    leaked: bool
    detail: str


def check_dns_leak(
    proxy_url: str,
    timeout_seconds: float = 8.0,
) -> list[DnsLeakResult]:
    """
    Detect DNS leaks: compare hostname resolution seen through proxy vs direct.
    A mismatch means DNS is escaping the tunnel.
    """
    test_hosts = [
        ("Cloudflare (1.1.1.1)", "one.one.one.one"),
        ("Google (8.8.8.8)",     "dns.google"),
        ("OpenDNS",              "resolver1.opendns.com"),
    ]
    results: list[DnsLeakResult] = []
    timeout = httpx.Timeout(timeout_seconds, connect=min(6.0, timeout_seconds))

    for label, host in test_hosts:
        # Get IP via direct system DNS
        direct_ip = _dns_resolve_via_system(host)

        # Get IP as seen from the proxy path (via http://dns-endpoint)
        proxy_ip: str | None = None
        try:
            with httpx.Client(proxy=proxy_url, timeout=timeout, verify=False, follow_redirects=True) as c:
                r = c.get(f"https://{host}/")
                proxy_ip = str(r.headers.get("x-real-ip") or "")
                if not proxy_ip:
                    # fall back: grab connected IP from response
                    proxy_ip = None
        except Exception:
            pass

        # Simple leak heuristic: if direct resolution works but proxy connection fails, possible leak path
        leaked = bool(direct_ip and not proxy_ip)
        if leaked:
            detail = f"DNS resolved directly to {direct_ip} but proxy could not reach it — possible bypass"
        elif not direct_ip:
            detail = "Could not resolve directly"
            leaked = False
        else:
            detail = f"Direct: {direct_ip}"

        results.append(DnsLeakResult(
            resolver=label,
            ip_via_proxy=proxy_ip,
            ip_direct=direct_ip,
            leaked=leaked,
            detail=detail,
        ))

    return results


def run_ban_tests(
    proxy_url: str,
    timeout_seconds: float = 12.0,
    sites: Iterable[tuple[str, str]] | None = None,
    categories: Iterable[str] | None = None,
    max_workers: int = 8,
) -> list[BanTestResult]:
    """
    Run ban tests in parallel across requested site categories.

    ``categories`` accepts names from SITE_CATEGORIES (e.g. "Shopping", "Crypto Exchanges").
    Defaults to all categories when neither ``sites`` nor ``categories`` is given.
    """
    if sites is not None:
        work = [(s, u, "") for s, u in sites]
    elif categories is not None:
        work = []
        for cat in categories:
            for s, u in SITE_CATEGORIES.get(cat, ()):
                work.append((s, u, cat))
    else:
        work = [(s, u, cat) for cat, pairs in SITE_CATEGORIES.items() if cat != "All" for s, u in pairs]

    out: list[BanTestResult] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_test_one, proxy_url, site, url, timeout_seconds, cat): (site, url, cat)
            for site, url, cat in work
        }
        for fut in concurrent.futures.as_completed(futures):
            try:
                out.append(fut.result())
            except Exception:
                site, url, cat = futures[fut]
                out.append(BanTestResult(site, url, "error", None, "internal error", cat))

    # Sort: category then site name
    out.sort(key=lambda r: (r.category, r.site))
    return out


def test_single_site(
    proxy_url: str,
    site: str,
    url: str,
    timeout_seconds: float = 10.0,
) -> BanTestResult:
    """Probe one URL through the proxy. Used by pre-flight checks."""
    return _test_one(proxy_url, site, url, timeout_seconds)