dns_server.py

#!/usr/bin/env python3

from argparse import ArgumentParser
from asyncio import run, get_running_loop, sleep, create_task
from logging import getLogger
from pathlib import Path
from dnslib import DNSRecord, DNSHeader, DNSQuestion, RR, A


logger = getLogger(Path(__file__).with_suffix('').name)

log_format = '%(asctime)s %(name)-25s %(levelname)5s: %(message)s'


def main():
    p = ArgumentParser()
    p.add_argument('--bind-host', default='127.0.0.1')
    p.add_argument('--bind-port', '-p', type=int, default=53)
    args = p.parse_args()
    setup_logging()
    run(async_main(args))


async def async_main(args):
    loop = get_running_loop()
    # https://docs.python.org/3/library/asyncio-protocol.html#udp-echo-server
    transport, protocol = await loop.create_datagram_endpoint(
        NameServerProtocol, local_addr=(args.bind_host, args.bind_port))
    try:
        while True:
            await sleep(3600)
    finally:
        transport.close()


class NameServerProtocol:

    def connection_made(self, transport):
        self.transport = transport
        logger.debug('connection made (transport=%r)', self.transport)

    def datagram_received(self, data, addr):
        logger.debug('datagram received from %r: %r (transport=%r)', addr, data, self.transport)
        task = create_task(process_udp_packet(data))

        def task_cb(future):
            reply_data = future.result()
            if reply_data:
                logger.debug('Sending data to %r: %r', addr, reply_data)
                self.transport.sendto(reply_data, addr)

        task.add_done_callback(task_cb)


async def process_udp_packet(data):
    record = DNSRecord.parse(data)
    logger.debug('parsed record: %r', record)
    await sleep(3)
    response = record.reply()
    response.add_answer(RR("abc.com",rdata=A("1.2.3.4")))
    logger.debug('response: %r', response)
    return response.pack()


def setup_logging():
    from logging import basicConfig, DEBUG
    basicConfig(format=log_format, level=DEBUG)


if __name__ == '__main__':
    main()