from dns import rdatatype, resolver
from requests import Session, adapters
from urllib.parse import urlparse
from urllib3.util.connection import HAS_IPV6


DNS_RESOLVER = resolver.Resolver()
DNS_RESOLVER.cache = resolver.LRUCache()  # type: ignore


class DNSAdapter(adapters.HTTPAdapter):
    def __init__(self, nameservers):
        self.nameservers = nameservers
        super().__init__()

    def resolve(self, host, nameservers):
        DNS_RESOLVER.nameservers = nameservers

        if HAS_IPV6:
            try:
                answers_v6 = DNS_RESOLVER.query(host, rdatatype.AAAA)
                for rdata_v6 in answers_v6:
                    return f"[{str(rdata_v6)}]"
            except resolver.NoAnswer:
                pass

        answers_v4 = DNS_RESOLVER.query(host, rdatatype.A)
        for rdata_v4 in answers_v4:
            return str(rdata_v4)

    def send(self, request, **kwargs):
        connection_pool_kwargs = self.poolmanager.connection_pool_kw
        result = urlparse(request.url)
        resolved_ip = self.resolve(result.hostname, self.nameservers)
        request.url = request.url.replace(result.hostname, resolved_ip)
        request.headers["Host"] = result.hostname
        request.headers[
            "User-Agent"
        ] = "Googlebot/2.1 (+http://www.google.com/bot.html)"

        if result.scheme == "https":
            connection_pool_kwargs["server_hostname"] = result.hostname
            connection_pool_kwargs["assert_hostname"] = result.hostname

        return super().send(request, **kwargs)


requests = Session()
requests.mount("http://", DNSAdapter(["1.1.1.1"]))
requests.mount("https://", DNSAdapter(["1.1.1.1"]))