fix: pin validated dns for ssrf checks

nanobot/security/network.py

@@ -5,9 +5,11 @@ from __future__ import annotations
 import ipaddress
 import re
 import socket
-from contextlib import suppress
+from contextlib import contextmanager, suppress
 from urllib.parse import urlparse
 
+import httpx
+
 _BLOCKED_NETWORKS = [
     ipaddress.ip_network("0.0.0.0/8"),
     ipaddress.ip_network("10.0.0.0/8"),
@@ -58,7 +60,7 @@ def _is_private(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
     return any(normalized in net for net in _BLOCKED_NETWORKS)
 
 
-def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
+def resolve_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str, tuple[str, ...]]:
     """Validate a URL is safe to fetch: scheme, hostname, and resolved IPs.
 
     ``allow_loopback`` is intentionally narrow: it only permits literal
@@ -66,26 +68,27 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool
     loopback. It does not allow RFC1918, link-local, metadata, or public DNS
     names that happen to resolve to loopback.
 
-    Returns (ok, error_message).  When ok is True, error_message is empty.
+    Returns (ok, error_message, resolved_ips).  When ok is True,
+    resolved_ips contains the public IPs that were validated for this URL.
     """
     try:
         p = urlparse(url)
     except Exception as e:
-        return False, str(e)
+        return False, str(e), ()
 
     if p.scheme not in ("http", "https"):
-        return False, f"Only http/https allowed, got '{p.scheme or 'none'}'"
+        return False, f"Only http/https allowed, got '{p.scheme or 'none'}'", ()
     if not p.netloc:
-        return False, "Missing domain"
+        return False, "Missing domain", ()
 
     hostname = p.hostname
     if not hostname:
-        return False, "Missing hostname"
+        return False, "Missing hostname", ()
 
     try:
         infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
     except socket.gaierror:
-        return False, f"Cannot resolve hostname: {hostname}"
+        return False, f"Cannot resolve hostname: {hostname}", ()
 
     addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
     for info in infos:
@@ -95,12 +98,70 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool
             continue
         addrs.append(addr)
     if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
-        return True, ""
+        return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
     for addr in addrs:
         if _is_private(addr):
-            return False, f"Blocked: {hostname} resolves to private/internal address {addr}"
+            return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
 
-    return True, ""
+    return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
+
+
+def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
+    """Validate a URL is safe to fetch: scheme, hostname, and resolved IPs."""
+    ok, error, _ = resolve_url_target(url, allow_loopback=allow_loopback)
+    return ok, error
+
+
+@contextmanager
+def pin_resolved_url_dns(url: str, resolved_ips: tuple[str, ...]):
+    """Pin DNS lookups for the URL hostname to previously validated IPs."""
+    try:
+        hostname = urlparse(url).hostname
+    except Exception:
+        hostname = None
+    if not hostname or not resolved_ips:
+        yield
+        return
+
+    pinned_host = hostname.rstrip(".").lower()
+    original_getaddrinfo = socket.getaddrinfo
+
+    def _getaddrinfo(host, port, family=0, type=0, proto=0, flags=0):  # noqa: A002
+        if str(host).rstrip(".").lower() != pinned_host:
+            return original_getaddrinfo(host, port, family, type, proto, flags)
+        infos = []
+        for ip in resolved_ips:
+            addr = ipaddress.ip_address(ip)
+            addr_family = socket.AF_INET6 if addr.version == 6 else socket.AF_INET
+            if family not in (0, socket.AF_UNSPEC, addr_family):
+                continue
+            sockaddr = (ip, port or 0, 0, 0) if addr_family == socket.AF_INET6 else (ip, port or 0)
+            infos.append((addr_family, type or socket.SOCK_STREAM, proto, "", sockaddr))
+        return infos
+
+    socket.getaddrinfo = _getaddrinfo
+    try:
+        yield
+    finally:
+        socket.getaddrinfo = original_getaddrinfo
+
+
+class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
+    """HTTPX transport that pins each request to the IPs validated for its URL."""
+
+    def __init__(self) -> None:
+        self._inner = httpx.AsyncHTTPTransport()
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        url = str(request.url)
+        ok, error, resolved_ips = resolve_url_target(url)
+        if not ok:
+            raise httpx.RequestError(error, request=request)
+        with pin_resolved_url_dns(url, resolved_ips):
+            return await self._inner.handle_async_request(request)
+
+    async def aclose(self) -> None:
+        await self._inner.aclose()
 
 
 def validate_resolved_url(url: str) -> tuple[bool, str]: