dnsdb-ip

#!/usr/bin/env python3

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

import os
import sys
import argparse
import json
import socket
#from multiprocessing import Pool        # utilise subprocesses rather than threads
from multiprocessing.dummy import Pool  # utilise threads rather than subprocesses
import signal
import requests
from io import StringIO


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

global url


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

def dnsdb_lookup(address):
    headers = {'Accept': 'application/json'}
    response = requests.get("%s/rdata/ip/%s" % (url, address), headers=headers)
    if "no results found for query" in response.text:
        return
    if "Internal server error" in response.text:
        sys.stderr.write('Remote server returned an Internal Server Error!')
        return
    data = StringIO(response.text)
    for line in data:
        jsonobj = json.loads(line)
        try:
            host = str(jsonobj['rrname'])[:-1] # slice the trailing dot
            sys.stdout.write('%s,%s\n' % (address, host))
            sys.stdout.flush()
        except KeyError as e:
            sys.stderr.write('%s => KeyError: %s [%s]\n' % (address, e, line))


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


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

if __name__ == '__main__':
    desc = 'Obtain observed hostnames from DNSDB for the supplied IP addresses and output results in CSV format.'

    parser = argparse.ArgumentParser(description=desc)
    parser.add_argument('file',
                        nargs='?',
                        type=argparse.FileType('r'),
                        action='store',
                        help='file containing a list of IP addresses split by a newline, otherwise read from STDIN',
                        metavar='FILE',
                        default=sys.stdin)
    args = parser.parse_args()

    # Check for DNSDB env
    try:
        global url
        url = os.environ['DNSDB']
        if not url:
            sys.stderr.write("Error: DNSDB environment variable is empty, unable to obtain server url, please set accordingly.\n")
            exit(1)
    except KeyError:
        sys.stderr.write("Error: DNSDB environment variable not set, unable to obtain server url, please set accordingly.\n")
        exit(1)

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

    # remove duplicates and sort
    addresses = list(set(addresses))
    addresses = sorted(addresses, key=lambda item: socket.inet_aton(item))

    # https://www.binpress.com/tutorial/simple-python-parallelism/121
    pool = Pool(processes=2, initializer=initializer)
    try:
        pool.map(dnsdb_lookup, addresses)
        pool.close()
        pool.join()
    except KeyboardInterrupt:
        pool.terminate()
        pool.join()