polycule-network/dns.py

import time
import socket
from ipaddress import IPv4Address, IPv4Network
from dnslib import DNSRecord,RCODE,QTYPE
from dnslib.server import DNSServer,DNSHandler,BaseResolver,DNSLogger
from dnslib.label import DNSLabel
from dnslib.ranges import IP4

class ProxyResolver(BaseResolver):
    def __init__(self, config):
        self.timeout = config.dns_timeout
        self.config = config
    
    def resolve(self, request, handler):
        qname = DNSLabel(request.q.qname)
        dns = self.config.dns_server(str(qname))

        try:
            proxy_r = request.send(dns.address, dns.port, timeout=self.timeout)
            reply = DNSRecord.parse(proxy_r)
        except socket.timeout:
            reply = request.reply()
            reply.header.rcode = getattr(RCODE, 'NXDOMAIN')
            return reply

        if dns.network != None and dns.network != config.local_network:
            for rr in reply.rr:
                rr.rdata.data = IPv4Address(config.translate(str(rr.rdata), network)).packed

        reply.set_header_qa()
        return reply

def translate(ip, untranslated_range, translated_range):
    for (loc, trans) in zip(IPv4Network(untranslated_range), IPv4Network(translated_range)):
        if str(loc) == ip:
            return str(trans)
    return ip


def run(config):
    resolver = ProxyResolver(config)
    handler = DNSHandler
    logger = DNSLogger("+request,+reply,+truncated,+error,-recv,-send,-data")
    udp_server = DNSServer(
            resolver,
            port=config.listen_port,
            address=config.listen_address,
            logger=logger,
            handler=handler
    )
    udp_server.start_thread()

    while udp_server.isAlive():
        time.sleep(1)