dns-a

#!/usr/bin/env python3

###########
# IMPORTS #
###########

import sys
import argparse
import asyncio
import aiodns

###########
# CLASSES #
###########

# TaskPool class taken from https://github.com/cgarciae/pypeln/blob/0.3.3/pypeln/task/utils.py

class TaskPool(object):
    def __init__(self, workers):
        self.semaphore = asyncio.Semaphore(workers) if workers else None
        self.tasks = set()
        self.closed = False

    async def put(self, coro):

        if self.closed:
            raise RuntimeError("Trying put items into a closed TaskPool")

        if self.semaphore:
            await self.semaphore.acquire()

        task = asyncio.create_task(coro)
        self.tasks.add(task)
        task.add_done_callback(self.on_task_done)
        task.set_exception

    def on_task_done(self, task):
        task
        self.tasks.remove(task)

        if self.semaphore:
            self.semaphore.release()

    async def join(self):
        await asyncio.gather(*self.tasks)
        self.closed = True

    async def __aenter__(self):
        return self

    def __aexit__(self, exc_type, exc, tb):
        return self.join()

#############
# FUNCTIONS #
#############

async def fetch(name, resolver, addresses):
    try:
        results = await resolver.query(name, 'A')
        for a in results:
            if addresses is None or a.host in addresses:
                sys.stdout.write('%s,%s\n' % (a.host, name))
                sys.stdout.flush()
    except Exception as e:
        try:
            err_number = e.args[0]
            err_text = e.args[1]
        except IndexError:
            sys.stderr.write(f"{name} => Couldn't parse exception: {e}\n")
        else:
            if err_number == 4:
                #sys.stderr.write(f"{address} => No record found.\n")
                pass
            elif err_number == 12:
                # Timeout from DNS server
                sys.stderr.write(f"{name} => Request timed out.\n")
            elif err_number == 1:
                # Server answered with no data
                pass
            else:
                sys.stderr.write(f"{name} =>  Unexpected exception: {e}\n")

async def run(loop, limit, hostnames, addresses):
    resolver = aiodns.DNSResolver(loop=loop, rotate=True)
    async with TaskPool(limit) as tasks:
        for name in hostnames:
            await tasks.put(fetch(name, resolver, addresses))

########
# MAIN #
########

if __name__ == '__main__':
    desc = 'Perform DNS lookup on supplied hostnames and output results in CSV format.'

    parser = argparse.ArgumentParser(description=desc)
    parser.add_argument('-n', '--nameservers',
                        nargs='?',
                        type=argparse.FileType('r'),
                        action='store',
                        help='file containing a list of nameservers split by a newline, otherwise use system resolver',
                        metavar='FILE',
                        default=None)
    parser.add_argument('-f', '--filter',
                        nargs='?',
                        type=argparse.FileType('r'),
                        action='store',
                        help='a list of IP addresses to filter against, only matching results are printed',
                        metavar='FILE',
                        default=None)
    parser.add_argument('file',
                        nargs='?',
                        type=argparse.FileType('r'),
                        action='store',
                        help='file containing a list of hostnames split by a newline, otherwise read from STDIN',
                        metavar='FILE',
                        default=sys.stdin)
    parser.add_argument('-l', '--limit',
                        type=int,
                        action='store',
                        help='maximum number of queries to run asynchronosly (default: 100)',
                        metavar='INT',
                        default=100)
    args = parser.parse_args()

    if args.filter:
        addresses = [line.strip() for line in args.filter if len(line.strip())>0 and line[0] != '#']
    else:
        addresses = None

    try:
        hostnames = [line.strip() for line in args.file if len(line.strip())>0 and line[0] != '#']
    except KeyboardInterrupt:
        exit()

    # remove duplicates and sort
    hostnames = sorted(set(hostnames))

    if args.nameservers:
        try:
            nameservers = [line.strip() for line in args.nameservers if len(line.strip())>0 and line[0] != '#']
        except KeyboardInterrupt:
            exit()
        
    loop = asyncio.get_event_loop()

    try:
        loop.run_until_complete(run(loop, args.limit, hostnames, addresses))
    except KeyboardInterrupt:
        sys.stderr.write("\nCaught keyboard interrupt, cleaning up...\n")
        asyncio.gather(*asyncio.Task.all_tasks()).cancel()
        loop.stop()
    finally:
        loop.close()