seansummers
6/28/2016 - 11:18 PM

DNS dict. __getitem__ any host, and it returns a randomized ipaddress from system DNS, honoring TTL

DNS dict. getitem any host, and it returns a randomized ipaddress from system DNS, honoring TTL

import random
import time

# if python2, install these batteries first
# ipaddress
import ipaddress
# dnspython
import dns.resolver


NX_EXPIRE = 5 * 60


def dns_query(host):
    try:
        query = dns.resolver.query(host)
        expiration = query.expiration
        addresses = tuple(ipaddress.ip_address(__.address) for __ in query)
    except dns.resolver.NXDOMAIN:
        expiration = time.time() + NX_EXPIRE
        addresses = ()
    return int(expiration), addresses


class DNS(dict):
    def __init__(self, hosts=None):
        hosts = hosts if hosts else ()
        data = {host: dns_query(host) for host in hosts}
        dict.__init__(self, data)

    def __getitem__(self, host):
        try:
            expiration, addresses = dict.__getitem__(self, host)
            if expiration <= int(time.time()):
                raise ValueError('DNS query expired')
        except (ValueError, KeyError):
            expiration, addresses = dns_query(host)
            self[host] = expiration, addresses
        address = random.choice(addresses) if addresses else None
        return address