Import source: axon/services/scraper

Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.

axon/services/scraper.py

@@ -0,0 +1,194 @@
+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