ImpTune/imptune/services/websearch.py

"""Web lookups: image search, driver-page search, and guarded image download.

The only place in ImpTune that talks to the internet. Everything here is
best-effort: a failed lookup returns an empty list or raises `WebSearchError`,
and the rest of the app keeps working offline.

DuckDuckGo is scraped rather than called through an API — no key, no signup.
That makes it fragile by nature: `search_images` / `search_pages` swallow parse
failures and return `[]` instead of a 500. Kill the whole feature with
`WEB_SEARCH=false`.

Fetching a URL the user picked means the *server* makes the request, so
`fetch_image` refuses anything that resolves to a private address. ImpTune runs
on the same LAN as the printers it configures — an unguarded fetcher would be a
port scanner for whoever can reach the UI.
"""
from __future__ import annotations

import html
import ipaddress
import json
import re
import socket
import urllib.error
import urllib.request
from dataclasses import dataclass
from urllib.parse import parse_qs, quote_plus, urlparse

import imptune.config as cfg

USER_AGENT = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/125.0 Safari/537.36"
)
TIMEOUT_SECONDS = 8
MAX_IMAGE_BYTES = 8 * 1024 * 1024
MAX_HTML_BYTES = 2 * 1024 * 1024
MAX_REDIRECTS = 3

DDG_HOME = "https://duckduckgo.com/"
DDG_IMAGES = "https://duckduckgo.com/i.js"
DDG_HTML = "https://html.duckduckgo.com/html/"

# Brands with a real vendor-wide driver. The search term is what a technician
# would type; the live results are whatever the engine returns for it.
GENERIC_DRIVER_TERMS: dict[str, str] = {
    "hp": "HP Universal Print Driver PCL6 download",
    "konica": "Konica Minolta Universal PCL Print Driver download",
    "konica minolta": "Konica Minolta Universal PCL Print Driver download",
    "xerox": "Xerox Global Print Driver download",
    "ricoh": "Ricoh PCL6 Universal Print Driver download",
    "lexmark": "Lexmark Universal Print Driver download",
    "brother": "Brother universal printer driver download",
    "canon": "Canon Generic PCL6 Printer Driver download",
    "epson": "Epson Universal Print Driver download",
    "kyocera": "Kyocera Classic Universal Print Driver download",
    "sharp": "Sharp Universal Print Driver download",
    "toshiba": "Toshiba Universal Printer 2 driver download",
    "oki": "OKI PCL universal print driver download",
    "samsung": "Samsung Universal Print Driver download",
    "dell": "Dell universal printer driver download",
}


# Display spellings for the search box's datalist, in the order a French shop
# is likeliest to need them.
BRAND_SUGGESTIONS = (
    "HP",
    "Konica Minolta",
    "Xerox",
    "Ricoh",
    "Canon",
    "Kyocera",
    "Brother",
    "Lexmark",
    "Epson",
    "Sharp",
    "Toshiba",
    "OKI",
    "Samsung",
    "Dell",
)


class WebSearchError(RuntimeError):
    """A lookup or download could not be completed."""


@dataclass(frozen=True)
class ImageHit:
    """One image search result."""

    url: str
    thumbnail: str
    title: str
    width: int
    height: int
    source: str


@dataclass(frozen=True)
class PageHit:
    """One web page search result."""

    url: str
    title: str
    snippet: str
    host: str


def enabled() -> bool:
    """False when the deployment opted out of outbound requests."""
    return bool(getattr(cfg, "WEB_SEARCH", True))


def _require_enabled() -> None:
    if not enabled():
        raise WebSearchError("Web lookups are disabled (WEB_SEARCH=false).")


# --------------------------------------------------------------------------
# SSRF guard
# --------------------------------------------------------------------------

def _is_public_ip(raw: str) -> bool:
    try:
        ip = ipaddress.ip_address(raw)
    except ValueError:
        return False
    return not (
        ip.is_private
        or ip.is_loopback
        or ip.is_link_local
        or ip.is_multicast
        or ip.is_reserved
        or ip.is_unspecified
    )


def assert_fetchable(url: str) -> None:
    """Raise `WebSearchError` unless `url` is an https(s) URL on a public host.

    Resolution happens here *and* the socket connects by hostname afterwards, so
    a DNS-rebinding host could still slip through between the two. Accepted:
    the payload is decoded by Pillow and discarded unless it is an image.
    """
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        raise WebSearchError("Only http and https URLs can be fetched.")
    if not parsed.hostname:
        raise WebSearchError("URL has no host.")

    try:
        infos = socket.getaddrinfo(parsed.hostname, None)
    except socket.gaierror as exc:
        raise WebSearchError(f"Cannot resolve {parsed.hostname}.") from exc

    for info in infos:
        if not _is_public_ip(info[4][0]):
            raise WebSearchError(
                f"{parsed.hostname} resolves to a private address — refused."
            )


class _GuardedRedirectHandler(urllib.request.HTTPRedirectHandler):
    """Re-run the SSRF guard on every redirect target."""

    def redirect_request(self, req, fp, code, msg, headers, newurl):  # noqa: D102
        assert_fetchable(newurl)
        return super().redirect_request(req, fp, code, msg, headers, newurl)


_opener = urllib.request.build_opener(_GuardedRedirectHandler)


def _get(
    url: str,
    *,
    referer: str | None = None,
    max_bytes: int,
    extra_headers: dict[str, str] | None = None,
) -> bytes:
    """GET `url` with the SSRF guard applied and the response body capped."""
    assert_fetchable(url)
    headers = {"User-Agent": USER_AGENT, "Accept-Language": "en-US,en;q=0.7"}
    if referer:
        headers["Referer"] = referer
    if extra_headers:
        headers.update(extra_headers)
    request = urllib.request.Request(url, headers=headers)
    try:
        with _opener.open(request, timeout=TIMEOUT_SECONDS) as response:
            return response.read(max_bytes + 1)[: max_bytes + 1]
    except urllib.error.URLError as exc:
        raise WebSearchError(f"Request failed: {exc}") from exc
    except (TimeoutError, socket.timeout) as exc:  # noqa: UP041 — socket.timeout on 3.9 paths
        raise WebSearchError("Request timed out.") from exc


# --------------------------------------------------------------------------
# Image search
# --------------------------------------------------------------------------

_VQD_RE = re.compile(rb"vqd=[\"']?([\w-]{8,})[\"']?")

# The image endpoint is an XHR route: without the fetch metadata a real browser
# would send, it answers 403 even with a valid vqd token. Verified header set —
# dropping any of these brings the 403 back.
_XHR_HEADERS = {
    "Accept": "application/json, text/javascript, */*; q=0.01",
    "X-Requested-With": "XMLHttpRequest",
    "Sec-Fetch-Dest": "empty",
    "Sec-Fetch-Mode": "cors",
    "Sec-Fetch-Site": "same-origin",
}


def _image_vqd(query: str) -> str:
    """Scrape the per-query token DuckDuckGo's image endpoint demands."""
    body = _get(
        f"{DDG_HOME}?q={quote_plus(query)}&iax=images&ia=images",
        max_bytes=MAX_HTML_BYTES,
    )
    match = _VQD_RE.search(body)
    if not match:
        raise WebSearchError("Could not read the search token.")
    return match.group(1).decode("ascii", "ignore")


def search_images(query: str, limit: int = 12) -> list[ImageHit]:
    """Image results for `query`, newest-first as the engine ranked them.

    Returns `[]` on any parse or transport failure — a dead scrape must not
    take a page down with it.
    """
    _require_enabled()
    query = query.strip()
    if not query:
        return []

    try:
        vqd = _image_vqd(query)
        raw = _get(
            f"{DDG_IMAGES}?l=us-en&o=json&q={quote_plus(query)}&vqd={vqd}&f=,,,&p=1",
            referer=DDG_HOME,
            max_bytes=MAX_HTML_BYTES,
            extra_headers=_XHR_HEADERS,
        )
        payload = json.loads(raw.decode("utf-8", "replace"))
    except (WebSearchError, json.JSONDecodeError):
        return []

    hits: list[ImageHit] = []
    for item in payload.get("results", []):
        url = item.get("image") or ""
        if not url.startswith(("http://", "https://")):
            continue
        hits.append(
            ImageHit(
                url=url,
                thumbnail=item.get("thumbnail") or url,
                title=(item.get("title") or "").strip(),
                width=int(item.get("width") or 0),
                height=int(item.get("height") or 0),
                source=urlparse(url).hostname or "",
            )
        )
        if len(hits) >= limit:
            break
    return hits


def fetch_image(url: str) -> bytes:
    """Download `url` and return its bytes (caller validates it is an image)."""
    _require_enabled()
    data = _get(url, max_bytes=MAX_IMAGE_BYTES)
    if len(data) > MAX_IMAGE_BYTES:
        raise WebSearchError("Image is larger than 8 MB.")
    if not data:
        raise WebSearchError("Server returned an empty response.")
    return data


# --------------------------------------------------------------------------
# Page search (driver downloads)
# --------------------------------------------------------------------------

_RESULT_RE = re.compile(
    r'<a[^>]+class="result__a"[^>]+href="(?P<href>[^"]+)"[^>]*>(?P<title>.*?)</a>',
    re.DOTALL,
)
_SNIPPET_RE = re.compile(
    r'<a[^>]+class="result__snippet"[^>]*>(?P<snippet>.*?)</a>', re.DOTALL
)
_TAG_RE = re.compile(r"<[^>]+>")


def _text(fragment: str) -> str:
    return html.unescape(_TAG_RE.sub("", fragment)).strip()


def _snippet_text(fragment: str) -> str:
    """Snippet text, minus the literal "undefined " DuckDuckGo's own template
    prefixes onto some results."""
    text = _text(fragment)
    return text[len("undefined ") :].lstrip() if text.startswith("undefined ") else text


def _unwrap(href: str) -> str:
    """Unwrap DuckDuckGo's `/l/?uddg=` click-tracking redirect."""
    if href.startswith("//"):
        href = "https:" + href
    parsed = urlparse(href)
    if "duckduckgo.com" in (parsed.hostname or "") and parsed.path.startswith("/l/"):
        target = parse_qs(parsed.query).get("uddg", [""])[0]
        if target:
            return target
    return href


def search_pages(query: str, limit: int = 8) -> list[PageHit]:
    """Web page results for `query`. `[]` on any failure, same as image search."""
    _require_enabled()
    query = query.strip()
    if not query:
        return []

    try:
        body = _get(
            f"{DDG_HTML}?q={quote_plus(query)}", max_bytes=MAX_HTML_BYTES
        ).decode("utf-8", "replace")
    except WebSearchError:
        return []

    snippets = [_snippet_text(m.group("snippet")) for m in _SNIPPET_RE.finditer(body)]
    hits: list[PageHit] = []
    for index, match in enumerate(_RESULT_RE.finditer(body)):
        url = _unwrap(html.unescape(match.group("href")))
        if not url.startswith(("http://", "https://")):
            continue
        hits.append(
            PageHit(
                url=url,
                title=_text(match.group("title")) or url,
                snippet=snippets[index] if index < len(snippets) else "",
                host=urlparse(url).hostname or "",
            )
        )
        if len(hits) >= limit:
            break
    return hits


# --------------------------------------------------------------------------
# Query builders
# --------------------------------------------------------------------------

def detect_brand(text: str) -> str | None:
    """Longest brand in `GENERIC_DRIVER_TERMS` that appears in `text`."""
    lowered = (text or "").lower()
    matches = [brand for brand in GENERIC_DRIVER_TERMS if brand in lowered]
    return max(matches, key=len) if matches else None


def generic_driver_query(text: str) -> str:
    """Search term for a vendor-wide driver, derived from a model name.

    A recognised brand gets that vendor's actual product name ("HP Universal
    Print Driver PCL6"); anything else falls back to the raw model plus
    "universal print driver download", which is what a technician would type.
    """
    text = (text or "").strip()
    brand = detect_brand(text)
    if brand:
        return GENERIC_DRIVER_TERMS[brand]
    return f"{text} universal print driver download".strip()


def image_query(text: str) -> str:
    """Default search term for a printer or driver icon."""
    text = (text or "").strip()
    if not text:
        return ""
    return f"{text} printer"