"""Refuse to fetch private / link-local addresses (SSRF guard).

Only ``read_page.py`` uses this today, and that asymmetry is deliberate:
screenshot and record hand back a *file path*, while web_read hands the page's
**text** straight to a model that is driven by user input. A prompt that says
"read http://169.254.169.254/latest/meta-data/iam/..." is an exfiltration
attempt, not a reading request.

``WEB_CAPTURE_ALLOW_PRIVATE_HOSTS=1`` opts out — needed when pointing the skill
at a dev server or an intranet page on purpose.

Caveat: the check resolves DNS once and then hands the URL to the browser,
which resolves it again. A name that flips between answers (DNS rebinding) can
slip through that gap. Closing it properly needs request-level interception in
Playwright; this guard is aimed at the ordinary case of a hostile or careless
URL, not at an attacker who controls a nameserver.
"""
from __future__ import annotations

import ipaddress
import os
import socket
from urllib.parse import urlparse

_TRUTHY = {"1", "true", "yes", "on"}


def allow_private_hosts() -> bool:
    return os.environ.get("WEB_CAPTURE_ALLOW_PRIVATE_HOSTS", "").strip().lower() in _TRUTHY


def _is_blocked(ip: ipaddress._BaseAddress) -> bool:
    return bool(
        ip.is_private
        or ip.is_loopback
        or ip.is_link_local
        or ip.is_reserved
        or ip.is_multicast
        or ip.is_unspecified
    )


def assert_public_url(raw_url: str) -> None:
    """Raise SystemExit when *raw_url* is not an ordinary public http(s) page."""
    parsed = urlparse(raw_url)
    if parsed.scheme not in ("http", "https"):
        raise SystemExit(
            f"--url 只接受 http/https，收到 {parsed.scheme or '(空)'}：{raw_url}"
        )
    host = parsed.hostname
    if not host:
        raise SystemExit(f"--url 缺少主机名：{raw_url}")

    if allow_private_hosts():
        return

    # A literal IP needs no DNS round-trip.
    try:
        literal = ipaddress.ip_address(host)
    except ValueError:
        literal = None
    if literal is not None:
        if _is_blocked(literal):
            raise SystemExit(_refusal(host, str(literal)))
        return

    try:
        infos = socket.getaddrinfo(host, parsed.port or (443 if parsed.scheme == "https" else 80))
    except socket.gaierror as e:
        raise SystemExit(f"无法解析主机 {host}：{e}")

    for info in infos:
        addr = info[4][0]
        try:
            ip = ipaddress.ip_address(addr)
        except ValueError:
            continue
        # Any resolved address being internal is enough to refuse — a name that
        # answers with one public and one private A record is the exact shape of
        # an SSRF attempt.
        if _is_blocked(ip):
            raise SystemExit(_refusal(host, addr))


def _refusal(host: str, addr: str) -> str:
    return (
        f"拒绝读取内网地址：{host} 解析到 {addr}（私有 / 环回 / 链路本地）。"
        "web_read 会把页面正文交给模型，因此默认只读公网页面。"
        "确实要读内网页面时设 WEB_CAPTURE_ALLOW_PRIVATE_HOSTS=1。"
    )