Axon/axon/services/scraper.py

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