Agent/cyber_agent/core/validator_web.py

"""Agentic Validator 的受控公开网页搜索与读取适配层。

本模块只提供两个只读能力：搜索公开页面和读取已授权 URL。它不复用
业务 Browser，不持有 Cookie，也不提供点击、登录、脚本执行或文件下载。
"""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import datetime, timezone
from hashlib import sha256
from html.parser import HTMLParser
import ipaddress
import json
import math
import os
import socket
from typing import Any, Protocol
from urllib.parse import urljoin, urlsplit, urlunsplit

import httpx

from cyber_agent.tools.registry import ToolRegistry


MAX_SEARCH_RESULTS = 5
MAX_REDIRECTS = 3
MAX_RESPONSE_BYTES = 512 * 1024
MAX_PAGE_CHARS = 16_000
MAX_TRACE_PAGE_CHARS = 50_000


class ValidatorWebError(RuntimeError):
    """Validator 网页工具的可呈现错误。"""


class ProviderUnavailable(ValidatorWebError):
    """搜索 Provider未配置、无密钥或暂不可用。"""


class UnsafeUrlError(ValidatorWebError):
    """URL、DNS或重定向越过公开网络只读边界。"""


class ValidatorWebSearchProvider(Protocol):
    """搜索服务的可注入端口；实现只返回普通网页候选。"""

    async def search(
        self,
        query: str,
        max_results: int,
    ) -> "ValidatorSearchResponse | Sequence[Mapping[str, Any]]":
        ...


@dataclass(frozen=True)
class ValidatorSearchResponse:
    """搜索适配器返回的结果和可选 Provider 成本。"""

    results: Sequence[Mapping[str, Any]]
    cost_usd: float = 0.0

    def __post_init__(self) -> None:
        if (
            isinstance(self.cost_usd, bool)
            or not isinstance(self.cost_usd, (int, float))
            or not math.isfinite(float(self.cost_usd))
            or self.cost_usd < 0
        ):
            raise ValueError("Validator search cost must be non-negative")


class SerperWebSearchProvider:
    """Serper Google Search API 的轻量适配器。"""

    endpoint = "https://google.serper.dev/search"

    def __init__(
        self,
        api_key: str | None = None,
        *,
        timeout_seconds: float = 15.0,
        cost_per_search_usd: float = 0.0,
        client_factory: Callable[[], httpx.AsyncClient] | None = None,
    ) -> None:
        self.api_key = (api_key or os.getenv("SERPER_API_KEY") or "").strip()
        self.timeout_seconds = timeout_seconds
        self.client_factory = client_factory
        self.cost_per_search_usd = ValidatorSearchResponse(
            results=[],
            cost_usd=cost_per_search_usd,
        ).cost_usd

    async def search(self, query: str, max_results: int) -> ValidatorSearchResponse:
        if not self.api_key:
            raise ProviderUnavailable("SERPER_API_KEY is not configured")
        client = (
            self.client_factory()
            if self.client_factory
            else httpx.AsyncClient(
                timeout=self.timeout_seconds,
                trust_env=False,
                follow_redirects=False,
            )
        )
        try:
            async with client:
                async with client.stream(
                    "POST",
                    self.endpoint,
                    headers={
                        "X-API-KEY": self.api_key,
                        "Content-Type": "application/json",
                    },
                    json={"q": query, "num": max_results},
                ) as response:
                    response.raise_for_status()
                    chunks: list[bytes] = []
                    size = 0
                    async for chunk in response.aiter_bytes():
                        size += len(chunk)
                        if size > MAX_RESPONSE_BYTES:
                            raise ProviderUnavailable(
                                "Serper response exceeds the 512 KiB limit"
                            )
                        chunks.append(chunk)
                    payload = json.loads(b"".join(chunks))
        except ProviderUnavailable:
            raise
        except (httpx.HTTPError, ValueError) as exc:
            raise ProviderUnavailable(f"Serper search failed: {exc}") from exc
        organic = payload.get("organic", []) if isinstance(payload, dict) else []
        return ValidatorSearchResponse(
            results=[
                item for item in organic if isinstance(item, Mapping)
            ][:max_results],
            cost_usd=self.cost_per_search_usd,
        )


class _VisibleTextParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__(convert_charrefs=True)
        self.title_parts: list[str] = []
        self.text_parts: list[str] = []
        self._ignored_depth = 0
        self._in_title = False

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        del attrs
        lowered = tag.lower()
        if lowered in {"script", "style", "noscript", "svg"}:
            self._ignored_depth += 1
        if lowered == "title":
            self._in_title = True

    def handle_endtag(self, tag: str) -> None:
        lowered = tag.lower()
        if lowered in {"script", "style", "noscript", "svg"} and self._ignored_depth:
            self._ignored_depth -= 1
        if lowered == "title":
            self._in_title = False

    def handle_data(self, data: str) -> None:
        text = " ".join(data.split())
        if not text or self._ignored_depth:
            return
        self.text_parts.append(text)
        if self._in_title:
            self.title_parts.append(text)

    @property
    def title(self) -> str:
        return " ".join(self.title_parts).strip()

    @property
    def text(self) -> str:
        return "\n".join(self.text_parts).strip()


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _source_id(url: str) -> str:
    return "src_" + sha256(url.encode("utf-8")).hexdigest()[:24]


def _blocked_ip(address: str) -> bool:
    ip = ipaddress.ip_address(address)
    if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped:
        ip = ip.ipv4_mapped
    return any((
        ip.is_private,
        ip.is_loopback,
        ip.is_link_local,
        ip.is_multicast,
        ip.is_reserved,
        ip.is_unspecified,
    ))


def normalize_public_url(url: str) -> str:
    """校验静态 URL 结构；DNS和每次重定向由异步检查继续确认。"""
    if not isinstance(url, str) or not url.strip():
        raise UnsafeUrlError("URL must be a non-empty string")
    parsed = urlsplit(url.strip())
    if parsed.scheme.lower() not in {"http", "https"}:
        raise UnsafeUrlError("Only HTTP(S) URLs are allowed")
    if parsed.username or parsed.password:
        raise UnsafeUrlError("URLs with user information are not allowed")
    host = (parsed.hostname or "").rstrip(".").lower()
    if not host or host == "localhost" or host.endswith(".local"):
        raise UnsafeUrlError("Local host names are not allowed")
    try:
        port = parsed.port
    except ValueError as exc:
        raise UnsafeUrlError(f"Invalid URL port: {exc}") from exc
    if port not in {None, 80, 443}:
        raise UnsafeUrlError("Only ports 80 and 443 are allowed")
    try:
        if _blocked_ip(host):
            raise UnsafeUrlError("Private or special-purpose IPs are not allowed")
    except ValueError:
        pass
    netloc = host
    if port is not None:
        netloc = f"[{host}]:{port}" if ":" in host else f"{host}:{port}"
    return urlunsplit((parsed.scheme.lower(), netloc, parsed.path or "/", parsed.query, ""))


async def _default_resolve(host: str, port: int) -> Sequence[str]:
    loop = asyncio.get_running_loop()
    infos = await loop.getaddrinfo(
        host,
        port,
        family=socket.AF_UNSPEC,
        type=socket.SOCK_STREAM,
    )
    return list({item[4][0] for item in infos})


async def require_public_dns(
    url: str,
    resolver: Callable[[str, int], Awaitable[Sequence[str]]] | None = None,
) -> str:
    normalized, _host, _port, _addresses = await _resolve_public_url(
        url,
        resolver,
    )
    return normalized


async def _resolve_public_url(
    url: str,
    resolver: Callable[[str, int], Awaitable[Sequence[str]]] | None,
) -> tuple[str, str, int, list[str]]:
    """解析一次公网地址，并返回后续建立连接所需的冻结 IP。"""
    normalized = normalize_public_url(url)
    parsed = urlsplit(normalized)
    host = parsed.hostname or ""
    port = parsed.port or (443 if parsed.scheme == "https" else 80)
    try:
        raw_addresses = await (resolver or _default_resolve)(host, port)
    except (OSError, socket.gaierror) as exc:
        raise UnsafeUrlError(f"URL host could not be resolved: {exc}") from exc
    addresses = list(dict.fromkeys(raw_addresses))
    if not addresses:
        raise UnsafeUrlError("URL host resolved to no addresses")
    for address in addresses:
        if _blocked_ip(address):
            raise UnsafeUrlError(
                f"URL host resolves to a private or special-purpose IP: {address}"
            )
    return normalized, host, port, addresses


def _pinned_request_url(url: str, address: str) -> str:
    """把逻辑 URL 替换为已校验 IP，防止 DNS 检查后再绑定。"""
    parsed = urlsplit(url)
    host = f"[{address}]" if ":" in address else address
    if parsed.port is not None:
        host = f"{host}:{parsed.port}"
    return urlunsplit((parsed.scheme, host, parsed.path, parsed.query, ""))


def _host_header(url: str) -> str:
    parsed = urlsplit(url)
    host = parsed.hostname or ""
    rendered = f"[{host}]" if ":" in host else host
    if parsed.port is not None:
        rendered = f"{rendered}:{parsed.port}"
    return rendered


def _decode_page(body: bytes, content_type: str) -> tuple[str, str]:
    mime = content_type.split(";", 1)[0].strip().lower()
    allowed = {"text/html", "text/plain", "application/json", "text/json"}
    if mime not in allowed and not mime.endswith("+json"):
        raise ValidatorWebError(f"Unsupported page content type: {mime or 'unknown'}")
    text = body.decode("utf-8", errors="replace")
    if mime == "text/html":
        parser = _VisibleTextParser()
        parser.feed(text)
        return parser.title, parser.text
    if "json" in mime:
        try:
            parsed = json.loads(text)
            text = json.dumps(parsed, ensure_ascii=False, sort_keys=True)
        except ValueError:
            pass
    return "", text


async def fetch_public_page(
    url: str,
    *,
    resolver: Callable[[str, int], Awaitable[Sequence[str]]] | None = None,
    client_factory: Callable[[], httpx.AsyncClient] | None = None,
) -> dict[str, Any]:
    """逐跳校验并读取一个公开文本页面，响应正文严格限长。"""
    current = url
    client = (
        client_factory()
        if client_factory
        else httpx.AsyncClient(
            timeout=20.0,
            trust_env=False,
            follow_redirects=False,
            headers={"Accept": "text/html,text/plain,application/json"},
        )
    )
    async with client:
        for hop in range(MAX_REDIRECTS + 1):
            current, host, _port, addresses = await _resolve_public_url(
                current,
                resolver,
            )
            client.cookies.clear()
            request = client.build_request(
                "GET",
                _pinned_request_url(current, addresses[0]),
                headers={
                    "Accept": "text/html,text/plain,application/json",
                    "Connection": "close",
                    "Host": _host_header(current),
                },
            )
            if urlsplit(current).scheme == "https":
                request.extensions["sni_hostname"] = host
            response = await client.send(request, stream=True)
            try:
                if response.status_code in {301, 302, 303, 307, 308}:
                    location = response.headers.get("location")
                    if not location:
                        raise ValidatorWebError("Redirect response has no Location header")
                    if hop >= MAX_REDIRECTS:
                        raise ValidatorWebError("Too many redirects")
                    current = urljoin(current, location)
                    continue
                response.raise_for_status()
                chunks: list[bytes] = []
                size = 0
                async for chunk in response.aiter_bytes():
                    size += len(chunk)
                    if size > MAX_RESPONSE_BYTES:
                        raise ValidatorWebError("Page exceeds the 512 KiB response limit")
                    chunks.append(chunk)
                body = b"".join(chunks)
                title, text = _decode_page(
                    body,
                    response.headers.get("content-type", ""),
                )
                truncated = len(text) > MAX_PAGE_CHARS
                text = text[:MAX_PAGE_CHARS]
                final_url = current
                return {
                    "source_id": _source_id(final_url),
                    "url": url,
                    "final_url": final_url,
                    "title": title,
                    "content_type": response.headers.get("content-type", ""),
                    "retrieved_at": _utc_now(),
                    "content_sha256": sha256(body).hexdigest(),
                    "text": text,
                    "truncated": truncated,
                    "untrusted_material": True,
                }
            finally:
                await response.aclose()
    raise ValidatorWebError("Page could not be fetched")


UsageRecorder = Callable[[int, int, float], Awaitable[None]]
PageFetcher = Callable[..., Awaitable[dict[str, Any]]]


@dataclass
class ValidatorToolLimits:
    max_searches: int
    max_opens: int
    max_tool_calls: int


@dataclass
class ValidatorToolSession:
    """单个 Scope Validator Trace 的私有工具、白名单和调用账本。"""

    provider: ValidatorWebSearchProvider | None
    allowed_urls: set[str]
    limits: ValidatorToolLimits
    usage_recorder: UsageRecorder | None = None
    resolver: Callable[[str, int], Awaitable[Sequence[str]]] | None = None
    page_fetcher: PageFetcher = fetch_public_page
    search_calls: int = 0
    open_calls: int = 0
    tool_calls: int = 0
    page_chars: int = 0
    opened_source_ids: set[str] = field(default_factory=set)
    _search_urls: set[str] = field(default_factory=set)

    def registry(self) -> ToolRegistry:
        registry = ToolRegistry()

        async def validator_web_search(query: str, max_results: int = 5) -> dict[str, Any]:
            return await self.search(query, max_results)

        async def validator_open_url(url: str) -> dict[str, Any]:
            return await self.open(url)

        registry.register(validator_web_search)
        registry.register(validator_open_url)
        return registry

    async def _reserve_call(self) -> None:
        self.tool_calls += 1
        if self.tool_calls > self.limits.max_tool_calls:
            raise ValidatorWebError("Validator tool-call limit reached")
        if self.usage_recorder:
            await self.usage_recorder(1, 0, 0.0)

    async def _record_material(
        self,
        material_chars: int,
        provider_cost_usd: float = 0.0,
    ) -> None:
        if (material_chars or provider_cost_usd) and self.usage_recorder:
            await self.usage_recorder(0, material_chars, provider_cost_usd)

    async def search(self, query: str, max_results: int = 5) -> dict[str, Any]:
        query = query.strip()
        if not query:
            raise ValidatorWebError("Search query must not be empty")
        if isinstance(max_results, bool) or not 1 <= max_results <= MAX_SEARCH_RESULTS:
            raise ValidatorWebError("max_results must be between 1 and 5")
        if self.search_calls >= self.limits.max_searches:
            raise ValidatorWebError("Validator search limit reached")
        self.search_calls += 1
        await self._reserve_call()
        if self.provider is None:
            return {
                "status": "provider_unavailable",
                "query": query,
                "results": [],
            }
        try:
            provider_result = await self.provider.search(query, max_results)
        except ProviderUnavailable as exc:
            return {
                "status": "provider_unavailable",
                "query": query,
                "error": str(exc),
                "results": [],
            }
        if isinstance(provider_result, ValidatorSearchResponse):
            raw_results = provider_result.results
            provider_cost_usd = provider_result.cost_usd
        else:
            raw_results = provider_result
            provider_cost_usd = 0.0
        results = []
        for index, item in enumerate(raw_results[:max_results], start=1):
            raw_url = str(item.get("url") or item.get("link") or "")
            try:
                normalized_url = normalize_public_url(raw_url)
            except ValidatorWebError:
                continue
            self._search_urls.add(normalized_url)
            results.append({
                "source_id": _source_id(normalized_url),
                "title": str(item.get("title") or ""),
                "url": normalized_url,
                "snippet": str(item.get("snippet") or ""),
                "rank": index,
                "retrieved_at": _utc_now(),
            })
        payload = {"status": "completed", "query": query, "results": results}
        await self._record_material(
            len(json.dumps(payload, ensure_ascii=False)),
            provider_cost_usd,
        )
        return payload

    async def open(self, url: str) -> dict[str, Any]:
        if self.open_calls >= self.limits.max_opens:
            raise ValidatorWebError("Validator page-open limit reached")
        normalized = normalize_public_url(url)
        normalized_allowed = {
            normalize_public_url(item) for item in self.allowed_urls
        }
        if normalized not in normalized_allowed and normalized not in self._search_urls:
            raise ValidatorWebError(
                "URL was not declared by the TaskReport or returned by this search session"
            )
        self.open_calls += 1
        await self._reserve_call()
        page = await self.page_fetcher(normalized, resolver=self.resolver)
        chars = len(str(page.get("text") or ""))
        if self.page_chars + chars > MAX_TRACE_PAGE_CHARS:
            raise ValidatorWebError("Validator page-text limit reached")
        self.page_chars += chars
        await self._record_material(chars)
        self.opened_source_ids.add(str(page["source_id"]))
        return page