devplacepy/devplacepy/services/devii/fetch/controller.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import ipaddress
import json
import logging
import re
import socket
from typing import Any
from urllib.parse import urlparse

import httpx

from devplacepy import stealth
from devplacepy.net_guard import effective_address, is_blocked_address
from ..config import Settings
from ..errors import NetworkError, ToolInputError, UpstreamError
from ..text import html_to_text

logger = logging.getLogger("devii.fetch")

CHROME_VERSION = "131"
CHROME_VERSION_FULL = "131.0.0.0"
STEALTH_HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
        f"(KHTML, like Gecko) Chrome/{CHROME_VERSION_FULL} Safari/537.36"
    ),
    "Accept": (
        "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,"
        "image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7"
    ),
    "Accept-Language": "en-US,en;q=0.9",
    "sec-ch-ua": (
        f'"Google Chrome";v="{CHROME_VERSION}", "Chromium";v="{CHROME_VERSION}", '
        '"Not_A Brand";v="24"'
    ),
    "sec-ch-ua-mobile": "?0",
    "sec-ch-ua-platform": '"Linux"',
    "Sec-Fetch-Dest": "document",
    "Sec-Fetch-Mode": "navigate",
    "Sec-Fetch-Site": "none",
    "Sec-Fetch-User": "?1",
    "Upgrade-Insecure-Requests": "1",
    "Cache-Control": "max-age=0",
    "DNT": "1",
}
TITLE = re.compile(r"<title[^>]*>(.*?)</title>", re.IGNORECASE | re.DOTALL)
MIN_FETCH_CHARS = 1000
ALLOWED_METHODS = ("GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS")


class FetchController:
    def __init__(self, settings: Settings) -> None:
        self._settings = settings

    async def dispatch(self, name: str, arguments: dict[str, Any]) -> str:
        if name == "fetch_url":
            return await self._fetch_url(arguments)
        if name == "http_request":
            return await self._http_request(arguments)
        raise ToolInputError(f"Unknown fetch tool: {name}")

    def _optional_cap(self, requested: Any, content: str) -> str:
        if requested is None:
            return content
        try:
            cap = max(MIN_FETCH_CHARS, int(requested))
        except (TypeError, ValueError):
            return content
        return content[:cap]

    async def _fetch_url(self, arguments: dict[str, Any]) -> str:
        url = str(arguments.get("url", "")).strip()
        if not url:
            raise ToolInputError("fetch_url requires a url.")
        if "://" not in url:
            url = "https://" + url
        await self._guard(url)

        body, final_url, content_type, status = await self._download(url)

        if (
            "html" in content_type
            or "xml" in content_type
            or content_type.startswith("text/")
        ):
            title_match = TITLE.search(body)
            title = html_to_text(title_match.group(1)) if title_match else ""
            content = html_to_text(body)
        elif "json" in content_type:
            title = ""
            content = body
        else:
            title = ""
            content = f"(non-text response: {content_type})"

        return json.dumps(
            {
                "status": "success",
                "url": final_url,
                "http_status": status,
                "content_type": content_type,
                "title": title,
                "content": self._optional_cap(arguments.get("max_chars"), content),
            },
            ensure_ascii=False,
        )

    def _coerce_mapping(self, value: Any, label: str) -> dict[str, Any] | None:
        if value is None:
            return None
        if isinstance(value, str):
            stripped = value.strip()
            if not stripped:
                return None
            try:
                value = json.loads(stripped)
            except json.JSONDecodeError as exc:
                raise ToolInputError(f"{label} must be a JSON object.") from exc
        if not isinstance(value, dict):
            raise ToolInputError(f"{label} must be a JSON object.")
        return {str(key): val for key, val in value.items()}

    def _coerce_json(self, value: Any) -> Any:
        if isinstance(value, str):
            try:
                return json.loads(value)
            except json.JSONDecodeError as exc:
                raise ToolInputError("json must be valid JSON.") from exc
        return value

    async def _http_request(self, arguments: dict[str, Any]) -> str:
        url = str(arguments.get("url", "")).strip()
        if not url:
            raise ToolInputError("http_request requires a url.")
        if "://" not in url:
            url = "https://" + url
        method = (str(arguments.get("method", "")).strip() or "GET").upper()
        if method not in ALLOWED_METHODS:
            raise ToolInputError(
                f"Unsupported HTTP method '{method}'. Use one of {', '.join(ALLOWED_METHODS)}."
            )
        await self._guard(url)

        headers = self._coerce_mapping(arguments.get("headers"), "headers")
        json_body = (
            self._coerce_json(arguments["json"])
            if arguments.get("json") is not None
            else None
        )
        form = (
            self._coerce_mapping(arguments.get("form"), "form")
            if arguments.get("form") is not None
            else None
        )
        raw_body = arguments.get("body")
        content = None if raw_body is None else str(raw_body)

        body, final_url, content_type, status, response_headers = await self._stream(
            method,
            url,
            headers=headers,
            json_body=json_body,
            form=form,
            content=content,
            raise_5xx=False,
        )

        if "json" in content_type or content_type.startswith("text/"):
            content_text = body
        elif "html" in content_type or "xml" in content_type:
            content_text = html_to_text(body)
        elif not content_type:
            content_text = body
        else:
            content_text = f"(non-text response: {content_type})"

        return json.dumps(
            {
                "status": "success",
                "url": final_url,
                "method": method,
                "http_status": status,
                "content_type": content_type,
                "headers": response_headers,
                "content": self._optional_cap(arguments.get("max_chars"), content_text),
            },
            ensure_ascii=False,
        )

    async def _guard(self, url: str) -> None:
        parsed = urlparse(url)
        if parsed.scheme not in ("http", "https"):
            raise ToolInputError("Only http and https URLs can be fetched.")
        host = parsed.hostname
        if not host:
            raise ToolInputError("URL has no host.")
        if self._settings.fetch_allow_private:
            return
        try:
            infos = await asyncio.to_thread(socket.getaddrinfo, host, None)
        except socket.gaierror as exc:
            raise NetworkError(f"Could not resolve host: {host}", host=host) from exc
        for info in infos:
            address = ipaddress.ip_address(info[4][0])
            if is_blocked_address(address):
                raise ToolInputError(
                    f"Refusing to fetch a private or local address ({effective_address(address)}). "
                    "Set DEVII_FETCH_ALLOW_PRIVATE=1 to override."
                )

    async def _download(self, url: str) -> tuple[str, str, str, int]:
        body, final_url, content_type, status, _ = await self._stream(
            "GET", url, raise_5xx=True
        )
        return body, final_url, content_type, status

    async def _stream(
        self,
        method: str,
        url: str,
        *,
        headers: dict[str, Any] | None = None,
        json_body: Any = None,
        form: dict[str, Any] | None = None,
        content: str | None = None,
        raise_5xx: bool = False,
    ) -> tuple[str, str, str, int, dict[str, str]]:
        limit = self._settings.fetch_max_bytes
        merged_headers = dict(STEALTH_HEADERS)
        if headers:
            merged_headers.update({str(k): str(v) for k, v in headers.items()})
        request_kwargs: dict[str, Any] = {}
        if json_body is not None:
            request_kwargs["json"] = json_body
        elif form is not None:
            request_kwargs["data"] = form
        elif content is not None:
            request_kwargs["content"] = content
        try:
            async with stealth.stealth_async_client(
                headers=merged_headers,
                follow_redirects=True,
                timeout=self._settings.fetch_timeout_seconds,
            ) as client:
                async with client.stream(method, url, **request_kwargs) as response:
                    if raise_5xx and response.status_code >= 500:
                        raise UpstreamError(
                            f"Server returned {response.status_code}.",
                            status=response.status_code,
                            url=url,
                        )
                    chunks: list[bytes] = []
                    total = 0
                    async for chunk in response.aiter_bytes():
                        chunks.append(chunk)
                        total += len(chunk)
                        if total >= limit:
                            break
                    raw = b"".join(chunks)[:limit]
                    encoding = response.encoding or "utf-8"
                    try:
                        body = raw.decode(encoding, errors="replace")
                    except LookupError:
                        body = raw.decode("utf-8", errors="replace")
                    content_type = response.headers.get("content-type", "").lower()
                    response_headers = {
                        key: value for key, value in response.headers.items()
                    }
                    return (
                        body,
                        str(response.url),
                        content_type,
                        response.status_code,
                        response_headers,
                    )
        except httpx.TimeoutException as exc:
            raise NetworkError(f"Request timed out fetching {url}", url=url) from exc
        except httpx.HTTPError as exc:
            raise NetworkError(f"Could not fetch {url}: {exc}", url=url) from exc