Import source: axon/services/scraper
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
This commit is contained in:
arsam
@@ -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
|
||||
Reference in New Issue
Block a user