ai-kefu-admin/tests/test_safe_fetch.py

"""safe_fetch：连接层 SSRF 校验（spec §3.1——防 DNS rebinding）。"""
import inspect
import socket

import httpx
import pytest

from app.platform.site_import import safe_fetch


def _mock_client_factory(handler):
    """redirect 类测试的打桩口：换掉 _make_client，注入 MockTransport client。"""

    def factory(timeout: float = 10.0) -> httpx.AsyncClient:
        return httpx.AsyncClient(
            transport=httpx.MockTransport(handler), follow_redirects=False)

    return factory


class _FakeRawSocket:
    def setsockopt(self, *args, **kwargs):
        pass


class _FakeAnyioStream:
    """anyio 流替身：吐一条完整 HTTP/1.1 200 响应（经 AnyIOStream 包装喂 httpcore）。"""

    def __init__(self) -> None:
        self._buf = (
            b"HTTP/1.1 200 OK\r\n"
            b"Content-Length: 2\r\n"
            b"Connection: close\r\n"
            b"\r\n"
            b"ok"
        )
        self._raw_socket = _FakeRawSocket()

    async def receive(self, max_bytes: int) -> bytes:
        chunk, self._buf = self._buf[:max_bytes], self._buf[max_bytes:]
        return chunk

    async def send(self, item: bytes) -> None:
        pass

    async def aclose(self) -> None:
        pass

    def extra(self, *args, **kwargs):
        return None


class TestVersionPin:
    def test_httpx私有口存在(self):
        """httpx/httpcore 升级破坏私有 API 时立即红（spec 要求的版本钉位）——真断言，禁永真。"""
        import httpcore
        t = httpx.AsyncHTTPTransport()
        assert hasattr(t, "_pool")
        assert "network_backend" in inspect.signature(
            httpcore.AsyncConnectionPool.__init__).parameters

    def test_httpcore私有导出口存在(self):
        """safe_fetch 消费的 httpcore 私有符号（AnyIOStream 流包装 / 异常映射口）逐个在位。"""
        from httpcore._backends.anyio import AnyIOStream
        from httpcore._exceptions import ConnectError, ConnectTimeout, map_exceptions
        assert callable(map_exceptions)
        assert issubclass(ConnectTimeout, Exception)
        assert issubclass(ConnectError, Exception)


class TestIsBlockedIp:
    @pytest.mark.parametrize("ip", [
        "127.0.0.1", "10.1.2.3", "172.16.0.1", "172.31.255.255",
        "192.168.1.1", "169.254.1.1", "0.0.0.0", "100.64.0.1",
        "::1", "fd00::1", "fe80::1", "::ffff:127.0.0.1", "::ffff:10.0.0.1",
    ])
    def test_私网段全拒(self, ip):
        assert safe_fetch.is_blocked_ip(ip) is True

    @pytest.mark.parametrize("ip", [
        "::", "224.0.0.1", "240.0.0.1", "255.255.255.255",
        "64:ff9b::1.2.3.4", "198.18.0.1",
    ])
    def test_补充禁段全拒(self, ip):
        """v6 未指定/组播/保留广播/NAT64/基准段——补丁钉位（::/128 是 v6 版 0.0.0.0）。"""
        assert safe_fetch.is_blocked_ip(ip) is True

    @pytest.mark.parametrize("ip", ["8.8.8.8", "1.1.1.1", "47.254.176.10", "2606:4700::1"])
    def test_公网放行(self, ip):
        assert safe_fetch.is_blocked_ip(ip) is False

    def test_非法IP串视为拒(self):
        assert safe_fetch.is_blocked_ip("not-an-ip") is True


class TestValidatingBackend:
    async def test_connect_tcp_拒绝私网IP字面量(self):
        backend = safe_fetch.ValidatingBackend()
        with pytest.raises(safe_fetch.BlockedAddressError):
            await backend.connect_tcp("127.0.0.1", 80)
        with pytest.raises(safe_fetch.BlockedAddressError):
            await backend.connect_tcp("10.0.0.1", 443)
        # （连接层校验逻辑用 IP 字面量直接钉，不做真实域名 DNS 测试）

    async def test_非法端口拒绝(self):
        backend = safe_fetch.ValidatingBackend()
        with pytest.raises(safe_fetch.BlockedAddressError):
            await backend.connect_tcp("8.8.8.8", 8080)


class TestFetch:
    async def test_非法端口拒绝(self):
        with pytest.raises(safe_fetch.FetchError):
            await safe_fetch.fetch("http://example.com:8080/")

    async def test_私网URL拒绝(self):
        with pytest.raises(safe_fetch.FetchError):
            await safe_fetch.fetch("http://127.0.0.1/x")
        with pytest.raises(safe_fetch.FetchError):
            await safe_fetch.fetch("http://10.0.0.1/x")

    async def test_非http协议拒绝(self):
        with pytest.raises(safe_fetch.FetchError):
            await safe_fetch.fetch("ftp://example.com/x")
        with pytest.raises(safe_fetch.FetchError):
            await safe_fetch.fetch("file:///etc/passwd")

    async def test_ipv6未指定地址URL拒绝(self):
        """http://[::]/ —— v6 版 0.0.0.0，URL 层必须拒（urlsplit 剥括号得 "::"）。"""
        with pytest.raises(safe_fetch.FetchError):
            await safe_fetch.fetch("http://[::]/x")

    async def test_逐跳重定向每跳重校验(self, monkeypatch):
        """302 → Location 指向私网 IP → 必须以 FetchError 终结（安全关键路径）。"""
        served = []

        def handler(request: httpx.Request) -> httpx.Response:
            served.append(str(request.url))
            return httpx.Response(302, headers={"Location": "http://10.0.0.2/evil"})

        monkeypatch.setattr(safe_fetch, "_make_client", _mock_client_factory(handler))
        with pytest.raises(safe_fetch.FetchError):
            await safe_fetch.fetch("http://example.com/start")
        # 首跳（公网）发生、第二跳（私网 Location）在 URL 层即被拦，未发起
        assert served == ["http://example.com/start"]

    async def test_重定向协议切换拒绝(self, monkeypatch):
        """302 → Location 换协议（ftp://）→ 第二跳 URL 层 FetchError。"""
        served = []

        def handler(request: httpx.Request) -> httpx.Response:
            served.append(str(request.url))
            return httpx.Response(302, headers={"Location": "ftp://evil/"})

        monkeypatch.setattr(safe_fetch, "_make_client", _mock_client_factory(handler))
        with pytest.raises(safe_fetch.FetchError):
            await safe_fetch.fetch("http://example.com/start")
        assert served == ["http://example.com/start"]

    async def test_连接层私网解析拒绝_真实链路(self, monkeypatch):
        """DNS rebinding 端到端钉位：URL 层查不出的私网目标（域名解析返回 10.x），
        经真实 _make_client 的连接层必须拦下——只打桩 getaddrinfo，不打桩 connect_tcp，
        证明校验先于任何连接发生。"""
        async def fake_getaddrinfo(host, port, **kwargs):
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.66.0.9", port))]

        monkeypatch.setattr(safe_fetch.anyio, "getaddrinfo", fake_getaddrinfo)
        with pytest.raises(safe_fetch.FetchError) as exc_info:
            await safe_fetch.fetch("http://rebinding.example.com/x")
        assert isinstance(exc_info.value.__cause__, safe_fetch.BlockedAddressError)

    async def test_重定向上限5跳(self, monkeypatch):
        """>5 跳 302 链 → FetchError。"""
        calls = []

        def handler(request: httpx.Request) -> httpx.Response:
            calls.append(request.url.path)
            return httpx.Response(302, headers={"Location": "/next"})

        monkeypatch.setattr(safe_fetch, "_make_client", _mock_client_factory(handler))
        with pytest.raises(safe_fetch.FetchError):
            await safe_fetch.fetch("http://example.com/start")
        # 首跳 + 5 次跟随 = 6 个请求；第 6 个 Location 不再跟
        assert len(calls) == 6

    async def test_五跳内重定向正常返回(self, monkeypatch):
        """恰好 5 跳 302 + 终点 200 → 正常返回（上限边界正向钉位）。"""
        locations = iter(["/r1", "/r2", "/r3", "/r4", "/final"])
        n_redirects = 0

        def handler(request: httpx.Request) -> httpx.Response:
            nonlocal n_redirects
            loc = next(locations, None)
            if loc is None:
                return httpx.Response(200, text="ok")
            n_redirects += 1
            return httpx.Response(302, headers={"Location": loc})

        monkeypatch.setattr(safe_fetch, "_make_client", _mock_client_factory(handler))
        resp = await safe_fetch.fetch("http://example.com/start")
        assert resp.status_code == 200
        assert resp.text == "ok"
        assert n_redirects == 5

    async def test_请求真实走校验后端(self, monkeypatch):
        """_make_client 的请求确实路由经 ValidatingBackend（连接层钩子被调用）。"""
        events = []

        async def fake_getaddrinfo(host, port, **kwargs):
            events.append(("resolve", host))
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port))]

        async def fake_connect_tcp(remote_host, remote_port, **kwargs):
            events.append(("connect", remote_host, remote_port))
            return _FakeAnyioStream()

        monkeypatch.setattr(safe_fetch.anyio, "getaddrinfo", fake_getaddrinfo)
        monkeypatch.setattr(safe_fetch.anyio, "connect_tcp", fake_connect_tcp)

        resp = await safe_fetch.fetch("http://example.com/")
        assert resp.status_code == 200
        assert resp.text == "ok"
        assert ("resolve", "example.com") in events
        # 连接目标是已验 IP 字面量而非原 hostname（DNS rebinding TOCTOU 闭合点）
        assert ("connect", "93.184.216.34", 80) in events