123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
"""Handles internet search requests and page fetching for use as AI chat context."""
|
|
|
|
import ipaddress
|
|
import socket
|
|
from urllib.parse import urlparse
|
|
|
|
import requests
|
|
from bs4 import BeautifulSoup
|
|
from ddgs import DDGS
|
|
|
|
# Maximum characters extracted from a fetched page before truncating.
|
|
MAX_PAGE_CHARS = 3000
|
|
|
|
_HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; AICodeEditor/1.0)"}
|
|
|
|
|
|
class SearchManager:
|
|
"""Performs DuckDuckGo searches and fetches web pages for AI context injection.
|
|
|
|
All outbound requests are validated against an SSRF blocklist so that
|
|
localhost and private network addresses can never be reached.
|
|
"""
|
|
|
|
def perform_search(self, query: str, max_results: int = 5) -> list[dict]:
|
|
"""Execute a DuckDuckGo text search and return normalised results.
|
|
|
|
Args:
|
|
query: The search query string.
|
|
max_results: Maximum number of results to return.
|
|
|
|
Returns:
|
|
List of {"title": str, "url": str, "snippet": str} dicts,
|
|
or an empty list if the search fails.
|
|
"""
|
|
try:
|
|
with DDGS() as ddgs:
|
|
raw = list(ddgs.text(query, max_results=max_results))
|
|
return self.parse_results(raw)
|
|
except Exception:
|
|
return []
|
|
|
|
def parse_results(self, raw_results: list[dict]) -> list[dict]:
|
|
"""Normalise raw DDGS result dicts to a consistent {"title", "url", "snippet"} shape.
|
|
|
|
Args:
|
|
raw_results: List of raw dicts returned by ddgs.text().
|
|
|
|
Returns:
|
|
Normalised list of result dicts.
|
|
"""
|
|
results = []
|
|
for r in raw_results:
|
|
results.append({
|
|
"title": r.get("title", ""),
|
|
"url": r.get("href", r.get("url", "")),
|
|
"snippet": r.get("body", r.get("snippet", "")),
|
|
})
|
|
return results
|
|
|
|
def fetch_page(self, url: str) -> str:
|
|
"""Fetch a web page and return its plain text content, truncated to MAX_PAGE_CHARS.
|
|
|
|
Args:
|
|
url: The URL to fetch.
|
|
|
|
Returns:
|
|
Plain text extracted from the page, or an error message string.
|
|
|
|
Raises:
|
|
ValueError: if the URL fails the SSRF safety check.
|
|
"""
|
|
self._validate_url(url)
|
|
try:
|
|
response = requests.get(url, timeout=10, headers=_HEADERS)
|
|
response.raise_for_status()
|
|
|
|
soup = BeautifulSoup(response.text, "html.parser")
|
|
|
|
# Remove non-content elements before extracting text.
|
|
for tag in soup(["script", "style", "nav", "footer"]):
|
|
tag.decompose()
|
|
|
|
text = soup.get_text(separator="\n", strip=True)
|
|
|
|
if len(text) > MAX_PAGE_CHARS:
|
|
text = text[:MAX_PAGE_CHARS] + "\n... [truncated]"
|
|
|
|
return text
|
|
|
|
except ValueError:
|
|
raise
|
|
except Exception as e:
|
|
return f"Error fetching page: {e}"
|
|
|
|
def _validate_url(self, url: str) -> None:
|
|
"""Block localhost, private IPs, and non-http(s) schemes to prevent SSRF attacks.
|
|
|
|
Args:
|
|
url: The URL to validate.
|
|
|
|
Raises:
|
|
ValueError: if the URL is considered unsafe.
|
|
"""
|
|
parsed = urlparse(url)
|
|
|
|
if parsed.scheme not in ("http", "https"):
|
|
raise ValueError(f"Blocked: only http/https allowed, got '{parsed.scheme}'")
|
|
|
|
hostname = parsed.hostname or ""
|
|
|
|
if hostname.lower() in ("localhost", "127.0.0.1", "::1"):
|
|
raise ValueError("Blocked: localhost access denied")
|
|
|
|
try:
|
|
ip = ipaddress.ip_address(socket.gethostbyname(hostname))
|
|
if ip.is_private or ip.is_loopback or ip.is_link_local:
|
|
raise ValueError(f"Blocked: private/loopback IP denied ({ip})")
|
|
except (socket.gaierror, ValueError) as e:
|
|
# Re-raise our own ValueError; ignore DNS resolution failures
|
|
# (let requests handle unknown hostnames naturally).
|
|
if isinstance(e, ValueError):
|
|
raise
|