AV-Watch/avwatch/http.py

"""Polite HTTP client: per-domain throttling, on-disk caching, backoff.

Every outbound request in avwatch goes through here. The scanner hits a few
dozen vendor sites and a search engine on a schedule, so being well-behaved
is not optional -- DuckDuckGo starts returning empty HTTP 202 bodies after
about two rapid queries, and several vendor sites sit behind Cloudflare.
"""

from __future__ import annotations

import hashlib
import json
import logging
import random
import time
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import urlsplit

import requests

log = logging.getLogger(__name__)

BROWSER_UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"
)

# Domains that need a real browser UA or they 403/404. Vendor marketing sites
# are the usual offenders.
_NEEDS_BROWSER_UA = (
    "duckduckgo.com",
    "blackmagicdesign.com",
    "shure.com",
    "waves.com",
    "novastar.tech",
    "renewedvision.com",
    "resolume.com",
    "focusrite.com",
    "reddit.com",
)

# Minimum seconds between requests to the same host. Search engines get a
# much wider berth than APIs that publish a documented rate limit.
_DEFAULT_DELAY = 1.5
_HOST_DELAY = {
    "html.duckduckgo.com": 6.0,
    "duckduckgo.com": 6.0,
    "api.github.com": 1.0,
    "hn.algolia.com": 1.0,
    "oauth.reddit.com": 1.2,
    "www.blackmagicdesign.com": 2.0,
}

# Status codes worth retrying. 202 is DuckDuckGo's "you are going too fast"
# tell: it returns 202 with a zero-length body rather than a 429.
_RETRY_STATUS = {202, 429, 500, 502, 503, 504}


class FetchError(RuntimeError):
    """A fetch failed in a way the caller should surface, not swallow."""


@dataclass
class Response:
    url: str
    status: int
    text: str
    from_cache: bool = False

    @property
    def ok(self) -> bool:
        return 200 <= self.status < 300 and bool(self.text)

    def json(self):
        return json.loads(self.text)


class HttpClient:
    def __init__(
        self,
        cache_dir: Path,
        cache_ttl: int = 3600,
        timeout: int = 30,
        max_retries: int = 3,
        user_agent: str | None = None,
        offline: bool = False,
    ):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.cache_ttl = cache_ttl
        self.timeout = timeout
        self.max_retries = max_retries
        self.offline = offline
        self.default_ua = user_agent or "avwatch/1.0 (church AV update monitor)"
        self._last_hit: dict[str, float] = {}
        self._session = requests.Session()
        self.stats = {"requests": 0, "cache_hits": 0, "errors": 0}

    # -- internals ---------------------------------------------------------

    def _cache_path(self, url: str, body: str | None) -> Path:
        key = hashlib.sha256((url + (body or "")).encode("utf-8")).hexdigest()[:32]
        host = urlsplit(url).netloc.replace(":", "_") or "unknown"
        return self.cache_dir / host / f"{key}.json"

    def _read_cache(self, path: Path, ttl: int) -> Response | None:
        if not path.exists():
            return None
        age = time.time() - path.stat().st_mtime
        if age > ttl:
            return None
        try:
            blob = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, ValueError):
            return None
        return Response(blob["url"], blob["status"], blob["text"], from_cache=True)

    def _write_cache(self, path: Path, resp: Response) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        payload = {"url": resp.url, "status": resp.status, "text": resp.text}
        try:
            path.write_text(json.dumps(payload), encoding="utf-8")
        except OSError as exc:  # a full disk should not kill the scan
            log.warning("could not cache %s: %s", resp.url, exc)

    def _throttle(self, host: str) -> None:
        delay = _HOST_DELAY.get(host, _DEFAULT_DELAY)
        last = self._last_hit.get(host)
        if last is not None:
            wait = delay - (time.time() - last)
            if wait > 0:
                time.sleep(wait)
        self._last_hit[host] = time.time()

    def _headers(self, url: str, extra: dict | None) -> dict:
        host = urlsplit(url).netloc
        ua = self.default_ua
        if any(host.endswith(d) for d in _NEEDS_BROWSER_UA):
            ua = BROWSER_UA
        headers = {
            "User-Agent": ua,
            "Accept": "*/*",
            "Accept-Language": "en-US,en;q=0.9",
        }
        if extra:
            headers.update(extra)
        return headers

    # -- public API --------------------------------------------------------

    def get(
        self,
        url: str,
        headers: dict | None = None,
        ttl: int | None = None,
        allow_cache: bool = True,
    ) -> Response:
        """GET a URL through the cache and throttle. Never raises on HTTP
        status; check ``resp.ok``. Raises FetchError only when offline with a
        cold cache."""
        ttl = self.cache_ttl if ttl is None else ttl
        path = self._cache_path(url, None)

        if allow_cache:
            cached = self._read_cache(path, ttl)
            if cached is not None:
                self.stats["cache_hits"] += 1
                return cached

        if self.offline:
            stale = self._read_cache(path, ttl=10**9)
            if stale is not None:
                self.stats["cache_hits"] += 1
                return stale
            raise FetchError(f"offline and nothing cached for {url}")

        host = urlsplit(url).netloc
        last: Response | None = None

        for attempt in range(self.max_retries):
            self._throttle(host)
            self.stats["requests"] += 1
            try:
                r = self._session.get(
                    url, headers=self._headers(url, headers), timeout=self.timeout
                )
            except requests.RequestException as exc:
                self.stats["errors"] += 1
                log.warning("%s: %s (attempt %d)", url, exc, attempt + 1)
                last = Response(url, 0, "")
                time.sleep(2**attempt + random.random())
                continue

            last = Response(url, r.status_code, r.text)
            if r.status_code in _RETRY_STATUS:
                # Back off hard; these are almost always rate limiting.
                sleep_for = (2 ** (attempt + 1)) + random.uniform(0, 2)
                log.info("%s returned %d, backing off %.1fs", host, r.status_code, sleep_for)
                time.sleep(sleep_for)
                continue
            break

        assert last is not None
        if last.ok:
            self._write_cache(path, last)
        else:
            self.stats["errors"] += 1
        return last

    def post_form(self, url: str, data: dict, headers: dict | None = None,
                  auth: tuple[str, str] | None = None) -> Response:
        """Uncached POST -- used only for OAuth token exchange."""
        if self.offline:
            raise FetchError(f"offline; cannot POST {url}")
        host = urlsplit(url).netloc
        self._throttle(host)
        self.stats["requests"] += 1
        try:
            r = self._session.post(
                url, data=data, headers=self._headers(url, headers),
                auth=auth, timeout=self.timeout,
            )
        except requests.RequestException as exc:
            self.stats["errors"] += 1
            raise FetchError(f"POST {url} failed: {exc}") from exc
        return Response(url, r.status_code, r.text)