HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
1"""Fetch transport: DIRECT MODE first (httpx, conditional requests, per-domain rate limits, robots.txt, backoff), optional2ESCALATED MODE (headless browser → Scrapfly → Firecrawl) that is never required and never automatic unless the connector asks.34 direct → retry → [browser] → [scrapfly] → [firecrawl] → BlockedError (→ review queue by the caller)5"""6from __future__ import annotations78import asyncio9import hashlib10import ipaddress11import logging12import socket13import time14from dataclasses import dataclass, field15from datetime import UTC, datetime16from urllib.parse import urljoin, urlparse, urlunparse17from urllib.robotparser import RobotFileParser1819import httpx2021from aiatlas.config import settings2223log = logging.getLogger(__name__)2425TRANSIENT_STATUS = {408, 425, 429, 500, 502, 503, 504}26BLOCK_STATUS = {401, 403, 451, 999}27REDIRECT_STATUS = {301, 302, 303, 307, 308}28MAX_REDIRECTS = 52930# ---------------------------------------------------------------------------------------------- SSRF guard31# Follow-up targets come from crawled content (links in model cards, feeds, READMEs): every destination — and every redirect32# hop — is validated before a connection is opened. Nothing internal is ever fetched.33_BLOCKED_SUFFIXES = (".local", ".internal", ".localhost", ".localdomain", ".lan", ".home", ".corp", ".intranet")34_BLOCKED_HOSTS = {"localhost", "metadata.google.internal", "metadata", "instance-data", "kubernetes.default", "kubernetes.default.svc"}35_BLOCKED_NETWORKS = [ipaddress.ip_network(n) for n in (36 "0.0.0.0/8", # unspecified / "this network"37 "10.0.0.0/8", # RFC191838 "100.64.0.0/10", # CGNAT39 "127.0.0.0/8", # loopback40 "169.254.0.0/16", # link-local (cloud metadata endpoints 169.254.169.254 …)41 "172.16.0.0/12", # RFC191842 "192.0.0.0/24", # IETF protocol assignments43 "192.168.0.0/16", # RFC191844 "198.18.0.0/15", # benchmarking45 "240.0.0.0/4", # reserved + broadcast46 "::/128", # unspecified47 "::1/128", # loopback48 "fc00::/7", # ULA49 "fe80::/10", # link-local50 "::ffff:0:0/96", # IPv4-mapped (checked again as IPv4 below)51 "64:ff9b::/96", # NAT6452)]535455class BlockedDestination(ValueError):56 """The URL points at a private, local or otherwise non-public destination."""575859def _ip_blocked(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:60 if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None:61 ip = ip.ipv4_mapped62 if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_unspecified or ip.is_reserved or ip.is_multicast:63 return True64 return any(ip in net for net in _BLOCKED_NETWORKS)656667def validate_destination(url: str, *, resolved_ips: list[str] | None = None) -> None:68 """Raise `BlockedDestination` for non-http(s) schemes, local/internal hostnames and private, loopback, link-local, CGNAT,69 IPv6 ULA/link-local or unspecified addresses. `resolved_ips` lets callers (and tests) inject the DNS answer."""70 p = urlparse(url.strip())71 if p.scheme.lower() not in ("http", "https"):72 raise BlockedDestination(f"blocked destination: scheme {p.scheme!r} for {url}")73 host = (p.hostname or "").strip().lower().rstrip(".")74 if not host:75 raise BlockedDestination(f"blocked destination: no host in {url}")76 if p.username or p.password:77 raise BlockedDestination(f"blocked destination: credentials in URL {url}")78 if host in _BLOCKED_HOSTS or host.endswith(_BLOCKED_SUFFIXES) or "." not in host and not _is_ip_literal(host):79 raise BlockedDestination(f"blocked destination: local host {host!r}")80 literal = _parse_ip(host)81 if literal is not None:82 if _ip_blocked(literal):83 raise BlockedDestination(f"blocked destination: non-public address {host}")84 return85 ips = resolved_ips if resolved_ips is not None else _resolve(host, p.port or (443 if p.scheme.lower() == "https" else 80))86 for raw in ips:87 ip = _parse_ip(raw)88 if ip is not None and _ip_blocked(ip):89 raise BlockedDestination(f"blocked destination: {host} resolves to non-public address {raw}")909192def _is_ip_literal(host: str) -> bool:93 return _parse_ip(host) is not None949596def _parse_ip(host: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None:97 try:98 return ipaddress.ip_address(host.strip("[]").split("%")[0])99 except ValueError:100 return None101102103def _resolve(host: str, port: int) -> list[str]:104 try:105 infos = socket.getaddrinfo(host, port, proto=socket.IPPROTO_TCP)106 except socket.gaierror:107 return [] # unresolvable hosts fail later with a transport error, not as SSRF108 return sorted({info[4][0] for info in infos})109110111async def validate_destination_async(url: str) -> None:112 """Non-blocking wrapper (DNS resolution runs in a thread)."""113 p = urlparse(url.strip())114 host = (p.hostname or "").strip().lower()115 if not host or _parse_ip(host) is not None:116 validate_destination(url, resolved_ips=[])117 return118 validate_destination(url, resolved_ips=[]) # scheme / hostname rules first (no DNS)119 ips = await asyncio.to_thread(_resolve, host, p.port or (443 if p.scheme.lower() == "https" else 80))120 validate_destination(url, resolved_ips=ips)121122123class FetchError(Exception):124 def __init__(self, message: str, *, status: int | None = None, url: str = ""):125 super().__init__(message)126 self.status = status127 self.url = url128129130class BlockedError(FetchError):131 """Access denied by the origin after every allowed transport (never bypass: park in the review queue)."""132133134class NotModified(Exception):135 """HTTP 304 — content unchanged since our stored validators."""136137138@dataclass139class FetchResult:140 url: str141 final_url: str142 status: int143 headers: dict[str, str]144 content: bytes145 content_type: str146 fetched_at: datetime147 duration_ms: int148 transport: str = "direct"149 from_cache: bool = False150 sha256: str = field(init=False)151152 def __post_init__(self) -> None:153 self.sha256 = hashlib.sha256(self.content).hexdigest()154155 @property156 def text(self) -> str:157 enc = "utf-8"158 ct = self.content_type.lower()159 if "charset=" in ct:160 enc = ct.split("charset=", 1)[1].split(";")[0].strip().strip('"') or "utf-8"161 try:162 return self.content.decode(enc, errors="replace")163 except LookupError:164 return self.content.decode("utf-8", errors="replace")165166 @property167 def is_html(self) -> bool:168 return "html" in self.content_type or self.content[:256].lstrip().lower().startswith((b"<!doctype html", b"<html"))169170 @property171 def is_json(self) -> bool:172 return "json" in self.content_type or self.content[:1] in (b"{", b"[")173174 @property175 def is_xml(self) -> bool:176 return "xml" in self.content_type or self.content[:5] == b"<?xml"177178 @property179 def is_pdf(self) -> bool:180 return "pdf" in self.content_type or self.content[:4] == b"%PDF"181182 def json(self): # type: ignore[no-untyped-def]183 import orjson184185 return orjson.loads(self.content)186187188def canonicalize_url(url: str) -> str:189 """Stable URL key: lowercase scheme/host, no fragment, no tracking params, no trailing slash duplication."""190 p = urlparse(url.strip())191 query = "&".join(sorted(q for q in p.query.split("&") if q and not q.lower().startswith(192 ("utm_", "ref=", "ref_", "fbclid", "gclid", "mc_cid", "mc_eid", "_hs", "igshid", "source="))))193 path = p.path or "/"194 if len(path) > 1 and path.endswith("/") and not path.endswith("//"):195 path = path.rstrip("/") or "/"196 return urlunparse((p.scheme.lower() or "https", p.netloc.lower(), path, "", query, ""))197198199def domain_of(url: str) -> str:200 host = urlparse(url).netloc.lower()201 return host.removeprefix("www.")202203204class _RateLimiter:205 """Per-domain minimum spacing between requests (token bucket, in-process). Cluster-wide fairness comes from Redis locks206 at the connector level (one connector run at a time)."""207208 def __init__(self) -> None:209 self._next: dict[str, float] = {}210 self._locks: dict[str, asyncio.Lock] = {}211212 async def wait(self, domain: str, per_min: int) -> None:213 lock = self._locks.setdefault(domain, asyncio.Lock())214 async with lock:215 spacing = 60.0 / max(1, per_min)216 now = time.monotonic()217 ready = self._next.get(domain, 0.0)218 if ready > now:219 await asyncio.sleep(ready - now)220 now = time.monotonic()221 self._next[domain] = now + spacing222223224class _Robots:225 def __init__(self) -> None:226 self._cache: dict[str, tuple[float, RobotFileParser | None]] = {}227228 async def allowed(self, client: httpx.AsyncClient, url: str) -> bool:229 if not settings.respect_robots:230 return True231 p = urlparse(url)232 key = f"{p.scheme}://{p.netloc}"233 cached = self._cache.get(key)234 if cached is None or cached[0] < time.monotonic():235 rp: RobotFileParser | None = RobotFileParser()236 try:237 r = await client.get(f"{key}/robots.txt", timeout=15, follow_redirects=False) # host already validated by the caller238 if r.status_code == 200 and len(r.content) < 512 * 1024:239 rp.parse(r.text.splitlines()) # type: ignore[union-attr]240 else:241 rp = None # no robots or error → allowed242 except Exception: # noqa: BLE001243 rp = None244 self._cache[key] = (time.monotonic() + 6 * 3600, rp)245 cached = self._cache[key]246 rp = cached[1]247 if rp is None:248 return True249 try:250 return rp.can_fetch(settings.user_agent.split("/")[0], url) or rp.can_fetch("*", url)251 except Exception: # noqa: BLE001252 return True253254255class Fetcher:256 """One Fetcher per connector run. Direct mode by default; `escalate=True` enables the optional chain."""257258 def __init__(self, *, rate_per_min: int | None = None, robots: bool = True, timeout_s: float | None = None,259 headers: dict[str, str] | None = None, http2: bool = True):260 self.rate_per_min = rate_per_min or settings.default_rate_per_min261 self.robots = robots262 self.timeout_s = timeout_s or settings.http_timeout_s263 self.headers = {"User-Agent": settings.user_agent, "Accept": "text/html,application/xhtml+xml,application/xml,application/json,text/*;q=0.9,*/*;q=0.8",264 "Accept-Language": "en-US,en;q=0.8", **(headers or {})}265 self._client: httpx.AsyncClient | None = None266 self._http2 = http2267268 async def __aenter__(self) -> Fetcher:269 # redirects are followed manually in `get()` so that every hop goes through the SSRF guard270 self._client = httpx.AsyncClient(headers=self.headers, timeout=httpx.Timeout(self.timeout_s, connect=20), follow_redirects=False,271 http2=self._http2, limits=httpx.Limits(max_connections=16, max_keepalive_connections=8))272 return self273274 async def __aexit__(self, *exc: object) -> None:275 if self._client:276 await self._client.aclose()277 self._client = None278279 @property280 def client(self) -> httpx.AsyncClient:281 if self._client is None:282 raise RuntimeError("use `async with Fetcher() as f`")283 return self._client284285 async def get(self, url: str, *, etag: str | None = None, last_modified: str | None = None, min_bytes: int = 1,286 escalate: bool = False, retries: int = 2, accept: str | None = None, rate_per_min: int | None = None) -> FetchResult:287 """Direct fetch with conditional headers. Raises NotModified (304), FetchError, BlockedError.288 Every destination and every redirect hop (max 5) is validated against the SSRF guard before a connection is opened."""289 client = self.client290 try:291 await validate_destination_async(url)292 except BlockedDestination as exc:293 raise FetchError(f"blocked destination: {exc}", url=url) from exc294 if self.robots and not await _robots.allowed(client, url):295 raise BlockedError(f"robots.txt disallows {url}", status=None, url=url)296 headers: dict[str, str] = {}297 if etag:298 headers["If-None-Match"] = etag299 if last_modified:300 headers["If-Modified-Since"] = last_modified301 if accept:302 headers["Accept"] = accept303 last_exc: Exception | None = None304 current = url305 hops = 0306 for attempt in range(retries + 1):307 await _limiter.wait(domain_of(current), rate_per_min or self.rate_per_min)308 t0 = time.perf_counter()309 try:310 async with client.stream("GET", current, headers=headers) as r:311 if r.status_code == 304:312 raise NotModified()313 if r.status_code in REDIRECT_STATUS:314 location = r.headers.get("location")315 if not location:316 raise FetchError(f"http {r.status_code} without Location for {current}", status=r.status_code, url=url)317 hops += 1318 if hops > MAX_REDIRECTS:319 raise FetchError(f"too many redirects (> {MAX_REDIRECTS}) for {url}", status=r.status_code, url=url)320 nxt = urljoin(current, location)321 try:322 await validate_destination_async(nxt)323 except BlockedDestination as exc:324 raise FetchError(f"blocked destination: redirect {current} → {nxt}: {exc}", status=r.status_code, url=url) from exc325 if domain_of(nxt) != domain_of(current):326 headers.pop("If-None-Match", None) # validators belong to the original resource only327 headers.pop("If-Modified-Since", None)328 current = nxt329 # a redirect is not a retry: loop again without consuming the attempt budget330 return await self._follow(client, current, headers, url=url, hops=hops, min_bytes=min_bytes, escalate=escalate,331 retries=retries, rate_per_min=rate_per_min)332 if r.status_code in TRANSIENT_STATUS and attempt < retries:333 retry_after = r.headers.get("retry-after")334 delay = min(60.0, float(retry_after)) if retry_after and retry_after.isdigit() else 2.0 * (2 ** attempt)335 log.info("transient status, backing off", extra={"url": url, "status": r.status_code, "delay": delay})336 await asyncio.sleep(delay)337 continue338 if r.status_code in BLOCK_STATUS:339 if escalate:340 return await self._escalate(url, reason=f"http {r.status_code}")341 raise BlockedError(f"http {r.status_code} for {url}", status=r.status_code, url=url)342 if r.status_code >= 400:343 raise FetchError(f"http {r.status_code} for {url}", status=r.status_code, url=url)344 return await self._read(r, url=url, final_url=current, t0=t0, min_bytes=min_bytes, escalate=escalate)345 except NotModified:346 raise347 except (BlockedError, FetchError):348 raise349 except (httpx.TimeoutException, httpx.TransportError) as exc:350 last_exc = exc351 if attempt < retries:352 await asyncio.sleep(1.5 * (2 ** attempt))353 continue354 raise FetchError(f"{exc.__class__.__name__}: {exc}", url=url) from exc355 raise FetchError(f"fetch failed: {last_exc}", url=url)356357 async def _follow(self, client: httpx.AsyncClient, current: str, headers: dict[str, str], *, url: str, hops: int, min_bytes: int,358 escalate: bool, retries: int, rate_per_min: int | None) -> FetchResult:359 """Continue a redirect chain: each hop is validated, capped at MAX_REDIRECTS, rate-limited per domain."""360 while True:361 await _limiter.wait(domain_of(current), rate_per_min or self.rate_per_min)362 t0 = time.perf_counter()363 try:364 async with client.stream("GET", current, headers=headers) as r:365 if r.status_code in REDIRECT_STATUS:366 location = r.headers.get("location")367 if not location:368 raise FetchError(f"http {r.status_code} without Location for {current}", status=r.status_code, url=url)369 hops += 1370 if hops > MAX_REDIRECTS:371 raise FetchError(f"too many redirects (> {MAX_REDIRECTS}) for {url}", status=r.status_code, url=url)372 nxt = urljoin(current, location)373 try:374 await validate_destination_async(nxt)375 except BlockedDestination as exc:376 raise FetchError(f"blocked destination: redirect {current} → {nxt}: {exc}", status=r.status_code, url=url) from exc377 if domain_of(nxt) != domain_of(current):378 headers.pop("If-None-Match", None)379 headers.pop("If-Modified-Since", None)380 current = nxt381 continue382 if r.status_code == 304:383 raise NotModified()384 if r.status_code in BLOCK_STATUS:385 if escalate:386 return await self._escalate(url, reason=f"http {r.status_code}")387 raise BlockedError(f"http {r.status_code} for {current}", status=r.status_code, url=url)388 if r.status_code >= 400:389 raise FetchError(f"http {r.status_code} for {current}", status=r.status_code, url=url)390 return await self._read(r, url=url, final_url=current, t0=t0, min_bytes=min_bytes, escalate=escalate)391 except (httpx.TimeoutException, httpx.TransportError) as exc:392 raise FetchError(f"{exc.__class__.__name__}: {exc}", url=url) from exc393394 async def _read(self, r: httpx.Response, *, url: str, final_url: str, t0: float, min_bytes: int, escalate: bool) -> FetchResult:395 chunks: list[bytes] = []396 size = 0397 async for chunk in r.aiter_bytes():398 size += len(chunk)399 if size > settings.max_body_bytes:400 raise FetchError(f"body exceeds {settings.max_body_bytes} bytes", status=r.status_code, url=url)401 chunks.append(chunk)402 content = b"".join(chunks)403 if len(content) < min_bytes:404 raise FetchError(f"suspiciously short body ({len(content)} bytes) for {url}", status=r.status_code, url=url)405 res = FetchResult(url=url, final_url=final_url, status=r.status_code, headers={k.lower(): v for k, v in r.headers.items()},406 content=content, content_type=r.headers.get("content-type", "").lower(), fetched_at=datetime.now(UTC),407 duration_ms=int((time.perf_counter() - t0) * 1000), transport="direct")408 if escalate and res.is_html and _looks_like_challenge(content):409 return await self._escalate(url, reason="anti-bot challenge page")410 return res411412 # ---------------------------------------------------------------------------------------------- escalation (optional)413 async def _escalate(self, url: str, *, reason: str) -> FetchResult:414 log.warning("escalating fetch", extra={"url": url, "reason": reason})415 if settings.browser_enabled:416 try:417 return await self._browser(url)418 except Exception as exc: # noqa: BLE001419 log.info("browser transport failed", extra={"url": url, "error": str(exc)})420 if settings.scrapfly_api_key:421 try:422 return await self._scrapfly(url)423 except Exception as exc: # noqa: BLE001424 log.info("scrapfly transport failed", extra={"url": url, "error": str(exc)})425 if settings.firecrawl_api_key:426 try:427 return await self._firecrawl(url)428 except Exception as exc: # noqa: BLE001429 log.info("firecrawl transport failed", extra={"url": url, "error": str(exc)})430 raise BlockedError(f"blocked after escalation ({reason}): {url}", url=url)431432 async def _browser(self, url: str) -> FetchResult:433 from playwright.async_api import async_playwright # optional dependency434435 t0 = time.perf_counter()436 async with async_playwright() as p:437 browser = await p.chromium.launch(headless=True)438 try:439 page = await browser.new_page(user_agent=settings.user_agent)440 resp = await page.goto(url, wait_until="networkidle", timeout=int(self.timeout_s * 1000))441 html = await page.content()442 status = resp.status if resp else 200443 finally:444 await browser.close()445 return FetchResult(url=url, final_url=url, status=status, headers={}, content=html.encode(), content_type="text/html; charset=utf-8",446 fetched_at=datetime.now(UTC), duration_ms=int((time.perf_counter() - t0) * 1000), transport="browser")447448 async def _scrapfly(self, url: str) -> FetchResult:449 t0 = time.perf_counter()450 r = await self.client.get("https://api.scrapfly.io/scrape", params={"key": settings.scrapfly_api_key, "url": url, "asp": "true",451 "render_js": "true", "country": "us"}, timeout=120)452 r.raise_for_status()453 data = r.json()["result"]454 content = (data.get("content") or "").encode()455 return FetchResult(url=url, final_url=data.get("url") or url, status=int(data.get("status_code") or 200), headers={},456 content=content, content_type=(data.get("content_type") or "text/html").lower(), fetched_at=datetime.now(UTC),457 duration_ms=int((time.perf_counter() - t0) * 1000), transport="scrapfly")458459 async def _firecrawl(self, url: str) -> FetchResult:460 t0 = time.perf_counter()461 r = await self.client.post("https://api.firecrawl.dev/v1/scrape", json={"url": url, "formats": ["rawHtml"]},462 headers={"Authorization": f"Bearer {settings.firecrawl_api_key}"}, timeout=120)463 r.raise_for_status()464 data = r.json().get("data", {})465 content = (data.get("rawHtml") or data.get("html") or "").encode()466 return FetchResult(url=url, final_url=(data.get("metadata") or {}).get("sourceURL") or url, status=200, headers={}, content=content,467 content_type="text/html", fetched_at=datetime.now(UTC), duration_ms=int((time.perf_counter() - t0) * 1000),468 transport="firecrawl")469470471def _looks_like_challenge(content: bytes) -> bool:472 head = content[:20000].lower()473 if len(content) > 60_000:474 return False475 markers = (b"just a moment", b"cf-chl-", b"challenge-platform", b"attention required", b"verify you are human", b"access denied",476 b"captcha", b"perimeterx", b"_px", b"datadome", b"enable javascript and cookies to continue")477 return sum(m in head for m in markers) >= 2478479480def file_result(path: str, *, url: str, content_type: str = "application/octet-stream") -> FetchResult:481 """Wrap a local file as a FetchResult (fixtures, seed snapshots, `aia run x --file`)."""482 with open(path, "rb") as fh:483 content = fh.read()484 return FetchResult(url=url, final_url=url, status=200, headers={}, content=content, content_type=content_type,485 fetched_at=datetime.now(UTC), duration_ms=0, transport="file")486487488_limiter = _RateLimiter()489_robots = _Robots()490491__all__ = ["MAX_REDIRECTS", "BlockedDestination", "BlockedError", "FetchError", "FetchResult", "Fetcher", "NotModified", "canonicalize_url", "domain_of",492 "file_result", "validate_destination", "validate_destination_async"]493