proxy_pool.py

"""
proxy_pool.py
--------------
A pool of proxy URLs and the rules for moving between them.

Why this exists as its own module: `--proxy` was a single static string
applied once at browser launch and never changed. That is the shape of a
demo, not of the thing proxies are bought for — the reason to hold a pool is
to spread a run across exits and to leave an exit that has started getting
challenged. The README sent readers to buy proxies without showing the
pattern; this is that pattern.

Kept engine-agnostic and side-effect free (no browser, no network) so the
rotation rules are covered by the offline suite rather than only by a live
run.

**Rotating the IP alone is not enough, and this is the part that is easy to
get wrong.** Carrying the same browser session across two exits is itself a
contradiction: cookies a bot manager issued against IP A, replayed from
IP B, are a stronger signal than either address on its own. So a caller must
build a FRESH browser context (new cookie jar, new storage) for every exit
this pool hands out — see playwright_scraper.py, which tears the browser
down and relaunches rather than swapping the proxy under a live session.
"""

from __future__ import annotations

import logging
import random
import re
import string
from typing import List, Optional
from urllib.parse import urlparse, urlunparse

logger = logging.getLogger("proxy_pool")

# `per-run` keeps one exit for the whole run — the safest default, since a
# single session that changes address mid-flight is more suspicious than one
# that does not. `per-page` takes a new exit for every page, which is what
# spreads volume; it costs a browser relaunch per page (see the module
# docstring for why that cost is mandatory rather than incidental).
ROTATE_MODES = ("per-run", "per-page")

# Schemes Playwright's `proxy.server` accepts. socks5 carries no credentials
# there (Chromium does not support authenticated SOCKS), so a socks5:// entry
# with a user:pass in it is rejected on load rather than silently ignored at
# request time.
_SUPPORTED_SCHEMES = ("http", "https", "socks5")


class ProxyError(ValueError):
    """A proxy list that cannot be used as given."""


def parse_proxy_line(line: str, source: str = "<arg>") -> Optional[str]:
    """Validate one proxy URL. Returns it, or None for a blank/comment line.

    Raises ProxyError with the offending line named, because a typo in a
    proxy list otherwise surfaces as a connection failure on page 1 with
    nothing pointing at the cause.
    """
    line = line.strip()
    if not line or line.startswith("#"):
        return None

    # Every message below reports mask(line), never the line itself. These
    # strings hold a password, and an error message is a log: a CI run
    # printed a live proxy login and password into its own public log this
    # way, from a traceback nobody expected to carry a credential.
    shown = mask(line)
    parsed = urlparse(line)
    if parsed.scheme not in _SUPPORTED_SCHEMES:
        raise ProxyError(
            f"{source}: {shown!r} — scheme must be one of "
            f"{', '.join(_SUPPORTED_SCHEMES)} (got {parsed.scheme or 'none'}). "
            f"A bare host:port is not enough; write http://host:port.")
    if not parsed.hostname:
        raise ProxyError(f"{source}: {shown!r} — no host in that URL.")
    # The PORT is validated here and not left to the first caller that reads
    # it. `urlparse` does not parse a port until you ask for one, and then it
    # raises ValueError — so a malformed entry sailed through this function
    # and blew up much later inside to_playwright as an uncaught traceback:
    # exit 1 (crash) where it should have been exit 2 (bad usage), with no
    # message saying what was wrong with the value.
    #
    # The value that caused it is worth knowing, because it is the mistake a
    # new user makes: a line from a proxy LIST FILE
    # ("http://host:port:login:password") pasted where a proxy URL belongs.
    # The extra colons become part of the port.
    try:
        parsed.port
    except ValueError:
        raise ProxyError(
            f"{source}: {shown!r} — the port is not a number. If you copied "
            f"this from a proxy list file, that format is "
            f"scheme://host:port:login:password and this expects a URL: "
            f"http://login:password@host:port") from None
    if parsed.scheme == "socks5" and (parsed.username or parsed.password):
        raise ProxyError(
            f"{source}: {shown!r} — Chromium cannot authenticate a SOCKS5 "
            f"proxy, so credentials here would be silently dropped. Use an "
            f"http:// entry for an authenticated proxy.")
    return line


def load_proxy_file(path: str) -> List[str]:
    """Read a proxy-per-line file. Blank lines and `#` comments are skipped."""
    proxies = []
    with open(path, "r", encoding="utf-8") as f:
        for lineno, raw in enumerate(f, 1):
            entry = parse_proxy_line(raw, source=f"{path}:{lineno}")
            if entry:
                proxies.append(entry)
    if not proxies:
        raise ProxyError(f"{path}: no proxy entries found (only blanks/comments?)")
    return proxies


def mask(url: Optional[str]) -> str:
    """A proxy URL safe to log: credentials replaced, host and port kept.

    Host and port stay visible on purpose — knowing WHICH exit a run used is
    the whole point of a rotation log, and it is not the secret.

    THIS FUNCTION MUST NEVER RAISE. It is the last thing standing between a
    password and a log, and it is called precisely when something is already
    wrong with the value. An earlier version read `parsed.port`, which
    `urlparse` computes lazily and which raises ValueError on a malformed
    authority — so the masker blew up on exactly the input that most needed
    masking, and the caller printed the raw string instead. That is how a
    live proxy login and password reached a public CI log.

    Anything it cannot take apart is redacted whole rather than echoed.
    """
    if not url:
        return "(none)"
    try:
        parsed = urlparse(url)
        host = parsed.hostname or "?"
        # `parsed.port` raises on a malformed authority; the raw netloc is
        # not safe to fall back to, because that is where the password is.
        try:
            port = f":{parsed.port}" if parsed.port else ""
        except ValueError:
            port = ":?"
        creds = "***:***@" if (parsed.username or parsed.password) else ""
        scheme = parsed.scheme or "?"
        # On a 2Captcha gateway every minted exit shares one host, one port
        # and one password -- the SESSION segment is the only thing telling
        # them apart. Without it three genuinely different exits log as three
        # identical lines, and a rotation log that cannot distinguish exits
        # is exactly the "looks like a pool, logs like a pool" failure this
        # feature has to avoid. The id is a per-run label, not the secret;
        # the password is still gone.
        tail = ""
        try:
            session = _SESSION_RE.search(parsed.username or "")
            if session and host.endswith(_GATEWAY_SUFFIX):
                tail = f" ({session.group(0).lstrip('-')})"
        except Exception:  # noqa: BLE001 -- never let the label break the mask
            tail = ""
        return f"{scheme}://{creds}{host}{port}{tail}"
    except Exception:  # noqa: BLE001 — a masker that raises is worse than a
        # vague one. Something is already wrong with this value; say so
        # without repeating it.
        return "(unparseable proxy URL, redacted)"


def to_playwright(url: Optional[str]) -> Optional[dict]:
    """Playwright's `proxy=` dict for a proxy URL, or None.

    Credentials go in their own fields rather than in `server`. Playwright
    passes `server` down to Chromium as a command-line switch, so a
    user:pass left in there would land in the browser process's argv — where
    anything on the machine that can run `ps` can read it.
    """
    if not url:
        return None
    parsed = urlparse(url)
    port = f":{parsed.port}" if parsed.port else ""
    proxy = {"server": f"{parsed.scheme}://{parsed.hostname}{port}"}
    if parsed.username:
        proxy["username"] = parsed.username
    if parsed.password:
        proxy["password"] = parsed.password
    return proxy


def split_credentials(url: Optional[str]):
    """(url_without_credentials, (username, password) or None).

    For the two engines that cannot take a proxy URL whole. Chromium's
    `--proxy-server=` switch has nowhere to put a password AND lands in the
    browser process's argv, where anything that can run `ps` reads it — so
    the address goes on the command line and the credentials go through the
    driver's own channel (pyppeteer's `page.authenticate`). Selenium has no
    such channel at all, which is why it strips these and warns.
    """
    if not url:
        return None, None
    parsed = urlparse(url)
    port = f":{parsed.port}" if parsed.port else ""
    scrubbed = f"{parsed.scheme}://{parsed.hostname}{port}"
    if parsed.username or parsed.password:
        return scrubbed, (parsed.username or "", parsed.password or "")
    return scrubbed, None


class ProxyPool:
    """An ordered pool of exits, plus a cursor and a rotation policy."""

    def __init__(self, proxies: List[str], rotate: str = "per-run",
                 shuffle: bool = False, rng: Optional[random.Random] = None):
        if not proxies:
            raise ProxyError("a proxy pool needs at least one entry")
        if rotate not in ROTATE_MODES:
            raise ProxyError(f"rotate must be one of {ROTATE_MODES}, got {rotate!r}")
        # Duplicates are dropped, order preserved. A pool is a set of EXITS,
        # and repeating one does not make it two: a list of fifty identical
        # entries — which is what a copied-and-pasted proxy list often is —
        # reported "exit 2/50" on every rotation while every one of them left
        # from the same address, and the single-exit warning below never
        # fired because it counted entries. A user then believes a run is
        # spread over fifty addresses when it is burning one.
        #
        # Said out loud rather than done silently: a pool quietly smaller
        # than the file that produced it is the same kind of surprise.
        seen, unique = set(), []
        for entry in proxies:
            if entry not in seen:
                seen.add(entry)
                unique.append(entry)
        if len(unique) < len(proxies):
            logger.warning(
                "Proxy pool: %d entries collapsed to %d distinct exit(s) — "
                "%d duplicate(s) dropped. Repeating an address does not "
                "spread a run across more of them.",
                len(proxies), len(unique), len(proxies) - len(unique))
        self._proxies = unique
        if shuffle:
            # Two runs started at the same minute otherwise hammer the same
            # first exit in the list.
            (rng or random).shuffle(self._proxies)
        self.rotate = rotate
        self._index = 0
        # Counted so a run can report how many exits it actually burned.
        self.rotations = 0

    def __len__(self) -> int:
        return len(self._proxies)

    @property
    def proxies(self) -> List[str]:
        """A copy of the exits, for handing a rotated view to each worker.

        A copy rather than the list itself: a worker builds its own pool from
        this, and two threads sharing one mutable list is the bug that makes
        concurrency stop being worth it.
        """
        return list(self._proxies)

    @property
    def current(self) -> str:
        return self._proxies[self._index % len(self._proxies)]

    def advance(self, reason: str) -> str:
        """Move to the next exit and return it. Wraps around the list.

        Wrapping rather than exhausting: a pool of 3 used across 50 pages is
        a legitimate configuration, and refusing to continue would be worse
        than reusing an exit. The log line says which exit and why, so a run
        that is cycling a too-small pool is visible rather than silent.
        """
        if len(self._proxies) == 1:
            logger.warning("Asked to rotate (%s) but the pool holds one exit "
                           "(%s) — staying on it. Add more with --proxy-file.",
                           reason, mask(self.current))
            return self.current
        self._index = (self._index + 1) % len(self._proxies)
        self.rotations += 1
        logger.info("Rotated proxy (%s) -> %s [exit %d/%d, rotation #%d]",
                    reason, mask(self.current), (self._index % len(self._proxies)) + 1,
                    len(self._proxies), self.rotations)
        return self.current

    def rotates_per_page(self) -> bool:
        return self.rotate == "per-page"


# ---------------------------------------------------------------------------
# Generating a pool from ONE credential (2Captcha's gateway convention)
# ---------------------------------------------------------------------------
# This is vendor-specific and deliberately kept here rather than applied to
# any `--proxy` URL. 2Captcha's residential gateway takes a `-session-{id}`
# segment inside the LOGIN and pins that session to one exit address for
# `-sessTime-{minutes}`; a different id is a different exit through the same
# gateway, host and password.
#
# MEASURED before this was written, because the whole idea rests on it:
# 2026-09-17, ten logins differing only in their session segment, one request
# each to an IP echo service -- ten answers, TEN DISTINCT ADDRESSES, no
# repeats. If the gateway had ignored the segment the result would have been
# twenty names for one IP: a pool that looks like a pool, logs like a pool,
# and quietly concentrates every request on one exit. That is worse than a
# file, which is why the check comes first and why the number is recorded.
_GATEWAY_SUFFIX = ".proxy.2captcha.com"
_SESSION_RE = re.compile(r"-session-[A-Za-z0-9]+")
_SESSTIME_RE = re.compile(r"-sessTime-\d+")
# Nine characters, matching the ids the vendor's own exported list uses.
_SESSION_CHARS = string.ascii_letters + string.digits
_SESSION_LEN = 9


def is_2captcha_gateway(url: Optional[str]) -> bool:
    """Is this one of 2Captcha's proxy gateways, where minting applies?"""
    if not url:
        return False
    try:
        host = (urlparse(url).hostname or "").lower()
    except ValueError:
        return False
    return host.endswith(_GATEWAY_SUFFIX)


def _session_id(rng: random.Random) -> str:
    return "".join(rng.choice(_SESSION_CHARS) for _ in range(_SESSION_LEN))


def with_session(url: str, session: str) -> str:
    """The same proxy URL with its login pinned to `session`.

    The rest of the login is left ALONE. A credential exported from the
    dashboard already carries `-zone-...-region-...` and often a session of
    its own; rebuilding the login from parts would drop whichever segment
    nobody thought of, so this replaces the session in place, or appends one
    before `-sessTime-` when there is none.
    """
    parts = urlparse(url)
    login = parts.username or ""
    if _SESSION_RE.search(login):
        new_login = _SESSION_RE.sub("-session-" + session, login, count=1)
    elif _SESSTIME_RE.search(login):
        new_login = _SESSTIME_RE.sub("-session-" + session + r"\g<0>", login, count=1)
    else:
        new_login = login + "-session-" + session
    password = parts.password or ""
    host = parts.hostname or ""
    port = f":{parts.port}" if parts.port else ""
    netloc = f"{new_login}:{password}@{host}{port}" if password else f"{new_login}@{host}{port}"
    return urlunparse((parts.scheme, netloc, parts.path, parts.params,
                       parts.query, parts.fragment))


def mint_sessions(url: str, count: int, rng: Optional[random.Random] = None) -> List[str]:
    """`count` copies of one gateway credential, each on its own session.

    Ids are unique within the run by construction: a collision would be two
    workers sharing an exit while the log claimed otherwise, which is the
    failure this whole idea has to avoid.
    """
    if count < 1:
        raise ProxyError("--proxy-sessions must be 1 or more")
    if not is_2captcha_gateway(url):
        raise ProxyError(
            "--proxy-sessions only applies to a 2Captcha proxy gateway "
            f"(a host under {_GATEWAY_SUFFIX}); the URL given is not one. "
            "Use --proxy-file for proxies from anywhere else.")
    rng = rng or random.Random()
    seen, out = set(), []
    while len(out) < count:
        sid = _session_id(rng)
        if sid in seen:
            continue
        seen.add(sid)
        out.append(with_session(url, sid))
    return out


def from_args(args) -> Optional[ProxyPool]:
    """Build a pool from --proxy-file / --proxy, or None if neither is set.

    `--proxy-file` wins when both are given, and says so: silently ignoring
    one of two conflicting options is how a run ends up on an exit the
    operator did not choose.
    """
    proxy_file = getattr(args, "proxy_file", None)
    single = getattr(args, "proxy", None)
    rotate = getattr(args, "proxy_rotate", "per-run")
    sessions = getattr(args, "proxy_sessions", None)

    if sessions and single and not proxy_file:
        proxies = mint_sessions(single, sessions)
        logger.info("Minted %d session-pinned exit(s) from one %s credential, "
                    "rotation: %s", len(proxies), mask(single), rotate)
        return ProxyPool(proxies, rotate=rotate,
                         shuffle=getattr(args, "proxy_shuffle", False))
    if sessions and proxy_file:
        logger.warning("--proxy-sessions is ignored when --proxy-file is given: "
                       "the file already names the exits.")

    if proxy_file:
        if single:
            logger.warning("--proxy-file and --proxy both given; using the file "
                           "and ignoring the single --proxy.")
        proxies = load_proxy_file(proxy_file)
        logger.info("Loaded %d proxy exit(s) from %s, rotation: %s",
                    len(proxies), proxy_file, rotate)
        return ProxyPool(proxies, rotate=rotate,
                         shuffle=getattr(args, "proxy_shuffle", False))

    if single:
        entry = parse_proxy_line(single, source="--proxy")
        if entry is None:
            raise ProxyError("--proxy was given but is empty")
        return ProxyPool([entry], rotate="per-run")

    return None