# retoor <retoor@molodetz.nl>
import asyncio
import ipaddress
import socket
from urllib.parse import urlsplit

import httpx

NAT64_PREFIX = ipaddress.ip_network("64:ff9b::/96")


class UnsafeURL(ValueError):
    pass


def _is_private(address):
    ip = ipaddress.ip_address(address)
    if isinstance(ip, ipaddress.IPv6Address):
        if ip.ipv4_mapped is not None:
            ip = ip.ipv4_mapped
        elif ip in NAT64_PREFIX:
            ip = ipaddress.IPv4Address(int(ip) & 0xFFFFFFFF)
    return (
        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 _resolve(host, port):
    return {info[4][0] for info in socket.getaddrinfo(host, port, proto=socket.IPPROTO_TCP)}


async def guard_public_host(host, port=443):
    if not host:
        raise UnsafeURL("missing host")
    try:
        addresses = await asyncio.to_thread(_resolve, host, port)
    except socket.gaierror as exc:
        raise UnsafeURL(f"cannot resolve {host}") from exc
    for address in addresses:
        if _is_private(address):
            raise UnsafeURL(f"{host} resolves to a non-public address")
    return addresses


async def guard_public_url(url):
    parts = urlsplit(url)
    if parts.scheme not in ("http", "https"):
        raise UnsafeURL("scheme must be http or https")
    port = parts.port or (443 if parts.scheme == "https" else 80)
    await guard_public_host(parts.hostname, port)
    return url


class GuardedTransport(httpx.AsyncBaseTransport):
    def __init__(self, inner):
        self.inner = inner

    async def handle_async_request(self, request):
        await guard_public_url(str(request.url))
        return await self.inner.handle_async_request(request)

    async def aclose(self):
        await self.inner.aclose()