import asyncio import ipaddress import logging import socket from urllib.parse import urldefrag, urljoin, urlparse import httpx import trafilatura from lxml import html from axon.config import settings logger = logging.getLogger(__name__) _HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; Axon/1.0; +https://github.com/axon)"} _ALLOWED_SCHEMES = {"http", "https"} def _is_safe_url(url: str) -> bool: """Return False for URLs that could be used for SSRF attacks.""" try: parsed = urlparse(url) if parsed.scheme not in _ALLOWED_SCHEMES: return False hostname = parsed.hostname if not hostname: return False # Resolve all IPs for the hostname and reject any private/internal address for _family, _type, _proto, _canonname, sockaddr in socket.getaddrinfo( hostname, None, proto=socket.IPPROTO_TCP ): addr = sockaddr[0] ip = ipaddress.ip_address(addr) if ( ip.is_loopback or ip.is_private or ip.is_link_local or ip.is_multicast or ip.is_reserved or ip.is_unspecified ): logger.warning("SSRF blocked: %s resolved to internal address %s", hostname, addr) return False return True except Exception as exc: logger.debug("URL safety check failed for %s: %s", url, exc) return False def normalize_url(url: str) -> str: clean, _fragment = urldefrag(url) return clean.rstrip("/") def same_site(url: str, start_url: str) -> bool: return urlparse(url).hostname == urlparse(start_url).hostname async def fetch_html(url: str, timeout: int = 10, retries: int = 2) -> str | None: if not await asyncio.to_thread(_is_safe_url, url): logger.warning("Skipping unsafe URL: %s", url) return None last_error: Exception | None = None for attempt in range(retries + 1): try: async with httpx.AsyncClient() as client: response = await client.get( url, timeout=timeout, follow_redirects=False, headers=_HEADERS, ) response.raise_for_status() return response.text except Exception as exc: last_error = exc if attempt < retries: await asyncio.sleep(0.25 * (attempt + 1)) logger.debug("Fetch failed for %s: %s", url, last_error) return None async def fetch_html_browser(url: str, timeout: int = 10) -> str | None: if not settings.enable_browser_fallback: return None if not await asyncio.to_thread(_is_safe_url, url): return None try: from playwright.async_api import async_playwright except Exception: logger.debug("Playwright is not installed; browser fallback disabled") return None try: async with async_playwright() as p: browser = await p.chromium.launch(headless=True) page = await browser.new_page() await page.goto(url, wait_until="networkidle", timeout=timeout * 1000) html_text = await page.content() await browser.close() return html_text except Exception as exc: logger.debug("Browser fetch failed for %s: %s", url, exc) return None async def extract_html(html_text: str, include_images: bool = False) -> tuple[str | None, str | None]: content = await asyncio.to_thread( trafilatura.extract, html_text, include_links=False, include_images=include_images, no_fallback=False, ) metadata = await asyncio.to_thread(trafilatura.extract_metadata, html_text) title = getattr(metadata, "title", None) if metadata else None return content, title async def extract_url( url: str, include_images: bool = False, timeout: int = 10, ) -> tuple[str | None, str | None, str | None]: html_text = await fetch_html(url, timeout=timeout) if html_text is None: html_text = await fetch_html_browser(url, timeout=timeout) if html_text is None: return None, None, "Fetch failed or unsafe URL" try: content, title = await extract_html(html_text, include_images=include_images) if not content: return "", title, "No extractable content found" return content, title, None except Exception as exc: logger.debug("Extraction failed for %s: %s", url, exc) return None, None, str(exc) async def discover_links(url: str, html_text: str, limit: int = 100) -> list[str]: def _parse_links() -> list[str]: try: doc = html.fromstring(html_text) except Exception: return [] links = [] for href in doc.xpath("//a[@href]/@href"): absolute = normalize_url(urljoin(url, href)) parsed = urlparse(absolute) if parsed.scheme in _ALLOWED_SCHEMES and same_site(absolute, url): links.append(absolute) return list(dict.fromkeys(links))[:limit] return await asyncio.to_thread(_parse_links) async def scrape_url(url: str, timeout: int = 5) -> str | None: try: html_text = await fetch_html(url, timeout=timeout) if html_text is None: html_text = await fetch_html_browser(url, timeout=timeout) if html_text is None: return None content, _title = await extract_html(html_text, include_images=False) return content except Exception as exc: logger.debug("Scrape failed for %s: %s", url, exc) return None async def scrape_urls( urls: list[str], max_concurrent: int = 10, timeout: int = 5, ) -> dict[str, str | None]: semaphore = asyncio.Semaphore(max_concurrent) async def _bounded(url: str) -> tuple[str, str | None]: async with semaphore: return url, await scrape_url(url, timeout=timeout) pairs = await asyncio.gather(*[_bounded(u) for u in urls], return_exceptions=True) result: dict[str, str | None] = {} for item in pairs: if isinstance(item, Exception): logger.debug("Scrape task exception: %s", item) continue url, content = item result[url] = content return result