dns-a

#!/usr/bin/env python2

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

import sys
import argparse
import socket
import dns
import dns.resolver
import dns.reversename
from multiprocessing import Pool
#from multiprocessing.dummy import Pool  # utilise threads rather than subprocesses
import signal


####################
# GLOBAL VARIABLES #
####################

global nameserver
global addresses


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

def resolve(hostname):
    max_attempts = 3
    resolver = dns.resolver.get_default_resolver()
    resolver.nameservers = [nameserver]
    resolver.lifetime = 3
    attempt = 0
    while attempt < max_attempts:
        try:
            answers = resolver.query(hostname)
        except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer, dns.resolver.NoNameservers):
            #sys.stderr.write('%s => No record found.\n' % (hostname))
            pass
        except dns.resolver.Timeout:
            attempt += 1
            if attempt == max_attempts:
                sys.stderr.write('%s => Request timed out.\n' % (hostname))
            continue
        else:
            for i in range(0, len(answers)):
                address = answers[i].address
                if addresses is None or address in addresses:
                    sys.stdout.write('%s,%s\n' % (address, hostname))
                    sys.stdout.flush()
        # break out of the loop
        attempt = max_attempts


def initializer():
    """Ignore CTRL+C in the worker process."""
    signal.signal(signal.SIGINT, signal.SIG_IGN)


########
# 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', '--nameserver',
                        action='store',
                        help='nameserver to use for lookups (default: 8.8.8.8)',
                        metavar='NS',
                        default='8.8.8.8')
    parser.add_argument('-f', '--filter',
                        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)
    args = parser.parse_args()

    global nameserver
    nameserver = args.nameserver

    global addresses
    if args.filter:
        with open(args.filter) as fp:
            addresses = [line.strip() for line in fp if len(line.strip())>0 and line[0] is not '#']
    else:
        addresses = None

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

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

    pool = Pool(processes=10, initializer=initializer)
    try:
        pool.map(resolve, hostnames)
        pool.close()
        pool.join()
    except KeyboardInterrupt:
        pool.terminate()
        pool.join()