aioproxies/src/aioproxies/manager.py

"""proxy source management: session templates, rotation, static.

`AioProxies` takes exactly one source (template / proxies / static) and hands out
`Proxy` objects via `next()`/`get()`. no module-level globals; a missing proxy file
raises, never `sys.exit`. see README for the full source/method reference.

proxy health (rotating list source only) is keyed by each proxy's canonical key
(`host:port:user:pass`, or `host:port` auth-less; port is normalized so `:080` and
`:80` collapse to one key). on template/static sources health/pool methods are
no-ops that log a warning and return, so generic caller code can call them
regardless of source.

per-proxy `timeout` (the availability field) is shared by cooldown and timed
burns: `None`/`0` fine, `-1` dead/permanent, a future unix ts = timed out until
then; durations only ever exist as call arguments, converted to `now + seconds`
and discarded, never stored raw.
"""
import logging
import random
import re
import string
import time
from typing import Dict, List, Optional, Union

from .proxy import Proxy, canonical_key, parse, to_proxy

log = logging.getLogger(__name__)

_DEAD = -1
_BARE_SESSION_SLOT = re.compile(r"(?<!\{)\{\}(?!\})")


class ProxiesExhaustedError(Exception):
    """raised by next() when every proxy in the pool is permanently dead (-1)"""


class AioProxies:
    """hands out proxies from a template, a rotating list, or a static value"""

    def __init__(
        self,
        *,
        template: Optional[str] = None,
        proxies: Optional[List[Union[str, Proxy]]] = None,
        static: Optional[Union[str, Proxy]] = None,
        session_len: int = 8,
        shuffle: bool = True,
        cooldown: int = 0,
    ):
        if proxies is not None and len(proxies) == 0:
            raise ValueError("proxies list is empty; provide at least one proxy")
        sources = [s for s in (template, proxies, static) if s is not None]
        if len(sources) != 1:
            raise ValueError("provide exactly one of: template, proxies, static")
        self.template = self._normalize_template(template) if template is not None else None
        self.session_len = session_len
        self.cooldown = cooldown
        self._shuffle = shuffle
        self._static = to_proxy(static) if static is not None else None
        self._proxies = self._dedupe([parse(p) for p in proxies]) if proxies is not None else []
        if self._proxies and shuffle:
            random.shuffle(self._proxies)
        self._state: Dict[str, Dict[str, object]] = {}
        for proxy in self._proxies:
            self._state.setdefault(proxy.key(), self._fresh_state())
        self._index = 0

    @staticmethod
    def _dedupe(proxies: List[Proxy]) -> List[Proxy]:
        """drop later entries sharing a canonical key, preserving first-seen order"""
        seen: Dict[str, Proxy] = {}
        for proxy in proxies:
            seen.setdefault(proxy.key(), proxy)
        return list(seen.values())

    @staticmethod
    def _fresh_state() -> Dict[str, object]:
        """a clean per-proxy state entry: no uses, no timeout, not burned"""
        return {"uses": 0, "timeout": None, "burned": False}

    @staticmethod
    def _normalize_template(template: str) -> str:
        """fill a bare `{}` session slot with `{session}`; leave escaped `{{}}` untouched"""
        return _BARE_SESSION_SLOT.sub("{session}", template)

    @classmethod
    def from_file(cls, path: str, **kwargs) -> "AioProxies":
        """build a rotating manager from a newline-delimited proxy file

        raises FileNotFoundError if the path is missing - never exits the process.
        """
        with open(path, "r", encoding="utf-8") as fh:
            lines = [ln.strip() for ln in fh if ln.strip()]
        if not lines:
            raise ValueError(f"no proxies found in {path}")
        return cls(proxies=lines, **kwargs)

    def session_id(self) -> str:
        """generate a fresh alphanumeric session id for template filling"""
        alphabet = string.ascii_letters + string.digits
        return "".join(random.choices(alphabet, k=self.session_len))

    def _is_list_source(self) -> bool:
        """whether this manager rotates a proxy list (vs template/static)"""
        return self._static is None and self.template is None

    def _warn_non_list(self, method: str) -> None:
        """log that a health/pool method only applies to list sources"""
        log.warning("%s applies only to list sources; no-op on template/static", method)

    def _available(self, timeout: object, now: float) -> bool:
        """whether a timeout value means the proxy is selectable right now"""
        if timeout is None or timeout == 0:
            return True
        if timeout == _DEAD:
            return False
        return now >= timeout

    @staticmethod
    def _clear_stale_burn(state: Dict[str, object]) -> None:
        """self-clear a lazily-expired timed burn's `burned` flag once found available

        mirrors the lazy expiry already applied to `timeout` itself: a timed burn()
        that has simply run out its clock should stop upgrading routine cooldown
        fallthrough logs to warning, without requiring an explicit restore().
        """
        if state["timeout"] not in (None, 0, _DEAD):
            state["burned"] = False

    def next(self, **fields: object) -> Proxy:
        """return the next proxy from the configured source

        template sources: `{session}` is always filled with a fresh id, other
        placeholders come from `fields` (e.g. next(country="ca", ttl=30)). fields
        are ignored by list/static sources.

        list sources: skips proxies whose timeout is active (-1 dead, or a future
        ts not yet passed), increments `uses` on handout, applies the manager's
        cooldown. if none are fine but some are merely timed, hands out the one
        recovering soonest (warns); raises ProxiesExhaustedError if all are dead.
        """
        if self._static is not None:
            return self._static
        if self.template is not None:
            if "session" in fields:
                # `session` is auto-filled with a fresh id; a caller-supplied one would
                # collide in str.format with an opaque TypeError - reject it clearly
                raise ValueError("'session' is filled automatically; do not pass it to next()")
            try:
                filled = self.template.format(session=self.session_id(), **fields)
            except (KeyError, IndexError) as exc:
                raise ValueError(
                    f"template placeholder {exc} not provided; pass it to next(**fields)"
                ) from exc
            except ValueError as exc:
                # a malformed template (e.g. an unmatched '{') makes str.format raise a
                # bare ValueError - re-raise naming the cause so it isn't cryptic
                raise ValueError(f"malformed proxy template {self.template!r}: {exc}") from exc
            return to_proxy(filled)
        return self._next_from_list()

    def _next_from_list(self) -> Proxy:
        """rotation + health selection over the proxy list"""
        if not self._proxies:
            raise ProxiesExhaustedError("no proxies in pool")

        now = time.time()
        count = len(self._proxies)
        fine: List[int] = []
        for offset in range(count):
            idx = (self._index + offset) % count
            state = self._state[self._proxies[idx].key()]
            if self._available(state["timeout"], now):
                fine.append(idx)
                self._clear_stale_burn(state)
                # only the first available proxy is used; scanning on is O(pool) per call
                break

        if fine:
            chosen = fine[0]
        else:
            chosen = self._soonest_recovering()
            if chosen is None:
                raise ProxiesExhaustedError("all proxies are permanently dead (-1)")
            if self._has_burned():
                log.warning("no proxies available; handing out the one recovering soonest")
            else:
                log.debug("all proxies cooling down; handing out the one recovering soonest")

        proxy = self._proxies[chosen]
        self._index = (chosen + 1) % count
        state = self._state[proxy.key()]
        state["uses"] = int(state["uses"]) + 1
        if self.cooldown > 0:
            # cooldown only ever EXTENDS availability delay - never shorten a longer
            # active timeout (a still-running timed burn) down to now+cooldown. _DEAD
            # (-1) is a permanent burn and must not be resurrected into a future ts.
            cooled = int(now) + self.cooldown
            current = state["timeout"]
            if current == _DEAD:
                pass
            elif isinstance(current, (int, float)) and current > cooled:
                pass
            else:
                state["timeout"] = cooled
        return proxy

    def _has_burned(self) -> bool:
        """whether any proxy carries a genuine burn (vs just cooldown resting), for log level choice"""
        return any(state["burned"] for state in self._state.values())

    def _soonest_recovering(self) -> Optional[int]:
        """index of the timed (non-dead) proxy recovering soonest, or None"""
        best_idx: Optional[int] = None
        best_ts: Optional[float] = None
        for idx, proxy in enumerate(self._proxies):
            timeout = self._state[proxy.key()]["timeout"]
            if timeout is None or timeout == 0 or timeout == _DEAD:
                continue
            ts = float(timeout)  # type: ignore[arg-type]
            if best_ts is None or ts < best_ts:
                best_ts = ts
                best_idx = idx
        return best_idx

    def get(self, **fields: object) -> Dict[str, str]:
        """convenience: next proxy as an aiohttp / aioweb proxies dict"""
        return self.next(**fields).aiohttp()

    def burn(self, proxy: Union[str, Proxy, Dict[str, str]], seconds: Optional[int] = None) -> None:
        """mark a proxy unusable: dead (-1) by default, or timed for `seconds`

        accepts any supported proxy shape. raises ValueError (naming the key) if the
        proxy is not in the pool. no-op + warning on template/static sources.
        """
        if not self._is_list_source():
            self._warn_non_list("burn()")
            return
        key = canonical_key(proxy)
        if key not in self._state:
            raise ValueError(f"proxy not in pool: {key}")
        if seconds is None:
            self._state[key]["timeout"] = _DEAD
        else:
            self._state[key]["timeout"] = int(time.time()) + seconds
        self._state[key]["burned"] = True

    def restore(self, proxy: Union[str, Proxy, Dict[str, str]]) -> None:
        """clear any burn/timeout on a proxy (back to fine). no-op if already fine

        matches remove()'s contract for an unknown proxy: logs a warning and returns
        rather than silently doing nothing.
        """
        if not self._is_list_source():
            self._warn_non_list("restore()")
            return
        key = canonical_key(proxy)
        if key not in self._state:
            log.warning("restore(): proxy not in pool: %s", key)
            return
        self._state[key]["timeout"] = None
        self._state[key]["burned"] = False

    def is_burned(self, proxy: Union[str, Proxy, Dict[str, str]]) -> bool:
        """whether a proxy is currently unavailable (lazy expiry of timed burns)

        -1 is always burned; an expired timed burn reads False. unknown proxies and
        non-list sources read False.
        """
        if not self._is_list_source():
            return False
        key = canonical_key(proxy)
        if key not in self._state:
            return False
        state = self._state[key]
        available = self._available(state["timeout"], time.time())
        if available:
            self._clear_stale_burn(state)
        return not available

    def stats(self) -> List[Dict[str, object]]:
        """per-proxy usage + state snapshot (list sources only, else empty)

        each entry: {"proxy": <spec>, "uses": int, "state": active|timed|dead,
        "timeout": <ts | -1 | None>}. the spec is the full canonical key.
        """
        if not self._is_list_source():
            self._warn_non_list("stats()")
            return []
        now = time.time()
        out: List[Dict[str, object]] = []
        for proxy in self._proxies:
            state = self._state[proxy.key()]
            timeout = state["timeout"]
            if timeout == _DEAD:
                label = "dead"
            elif self._available(timeout, now):
                label = "active"
            else:
                label = "timed"
            out.append({
                "proxy": proxy.key(),
                "uses": int(state["uses"]),
                "state": label,
                "timeout": timeout,
            })
        return out

    def reset_stats(self) -> None:
        """zero every proxy's use counter; leave timeouts untouched"""
        if not self._is_list_source():
            self._warn_non_list("reset_stats()")
            return
        for state in self._state.values():
            state["uses"] = 0

    def replace(self, proxies: List[Union[str, Proxy, Dict[str, str]]], *, keep_state: bool = False) -> None:
        """swap the entire proxy list

        keep_state=False (default) wipes all per-proxy state (fresh batch).
        keep_state=True preserves uses/timeout for proxies whose canonical key
        survives the swap; new proxies start clean, dropped ones are forgotten.
        resets the rotation index and honors the manager's shuffle setting. raises
        ValueError on an empty list, matching the constructor's guard - use remove()
        to drop proxies one at a time if you actually want a smaller pool.
        """
        if not self._is_list_source():
            self._warn_non_list("replace()")
            return
        if len(proxies) == 0:
            raise ValueError("proxies list is empty; provide at least one proxy")
        old_state = self._state
        incoming = [to_proxy(p) for p in proxies]
        if self._shuffle:
            random.shuffle(incoming)
        new_proxies: List[Proxy] = []
        new_state: Dict[str, Dict[str, object]] = {}
        for proxy in incoming:
            key = proxy.key()
            if key in new_state:
                continue
            new_proxies.append(proxy)
            if keep_state and key in old_state:
                new_state[key] = old_state[key]
            else:
                new_state[key] = self._fresh_state()
        self._proxies = new_proxies
        self._state = new_state
        self._index = 0

    def add(self, proxies: Union[str, Proxy, Dict[str, str], List[Union[str, Proxy, Dict[str, str]]]]) -> None:
        """append proxies to the pool, keeping existing state; skip duplicate keys"""
        if not self._is_list_source():
            self._warn_non_list("add()")
            return
        items = proxies if isinstance(proxies, list) else [proxies]
        for item in items:
            proxy = to_proxy(item)
            key = proxy.key()
            if key in self._state:
                continue
            self._proxies.append(proxy)
            self._state[key] = self._fresh_state()

    def remove(self, proxy: Union[str, Proxy, Dict[str, str]]) -> None:
        """drop a proxy from the pool entirely (by canonical key, any shape)

        distinct from burn (burn = unusable but tracked; remove = gone). decrements
        the rotation index when the removed slot precedes it, so the next() cursor
        still lands on the same upcoming proxy instead of skipping one. no-op +
        warning if the proxy is not present.
        """
        if not self._is_list_source():
            self._warn_non_list("remove()")
            return
        key = canonical_key(proxy)
        if key not in self._state:
            log.warning("remove(): proxy not in pool: %s", key)
            return
        removed_idx = next(i for i, p in enumerate(self._proxies) if p.key() == key)
        self._proxies = [p for p in self._proxies if p.key() != key]
        del self._state[key]
        if not self._proxies:
            self._index = 0
        else:
            if removed_idx < self._index:
                self._index -= 1
            self._index %= len(self._proxies)


# name aliases - same class, call it whichever reads best at your call site
ProxyManager = AioProxies
aioproxies = AioProxies