dns_resolver.py

import dns.resolver

from dns.exception import DNSException
from threading import Thread
from Queue import Queue


class DNSResolver(object):
    """
    Attributes:
        custom_configure: True if you want to specify your own nameservers. False by default.
        custom_nameservers: list your nameservers here if you want to use your
            own nameservers. Not does not take effect unless custom_configure is True
    """
    def __init__(self, custom_configure=False, custom_nameservers=[]):
        if custom_configure:
            if not isinstance(custom_nameservers, list):
                raise InputError(custom_nameservers, "custom_nameservers need to be a list")
            if len(custom_nameservers) == 0:
                raise InputError(custom_nameservers, "speciify nameservers to resolve from")

        """
        The configure flag is True for default configuration
        """
        self.resolver = dns.resolver.Resolver(configure=not custom_configure)

        if custom_configure:
            self.resolver.nameservers = custom_nameservers

        """
        This is related to the NS and MX query records that we do.
        We only go up max_depth of parents
        """
        self.max_depth = 2
        self.supported_records = ["NS", "A", "MX", "CNAME"]

    def getRecords(self, hostname, record):
        """
        getRecords should be able to get a record for a hostname
        Currently, the record could be of type listed in records.
        It is not clear that we need to restrict the queries for records
        that we allow to be made so might make sense to get rid of it in
        the future

        Attributes:
            hostname: hostname to query
            record: the type of record that we want to fetch
        """
        if record not in self.supported_records:
            raise InputError(record_type, "Unsupported record type: {}".format(record_type))

        records = self.resolver.query(hostname, record)
        return [rec.to_text() for rec in records]


    def getRecordsBrief(self, hostname, record_type):
        if record_type not in self.supported_records:
            raise InputError(record_type, "Unsupported record type: {}".format(record_type))

        try:
            records = []
            fetch_err = None
            records = self.getRecords(hostname, record_type)
        except DNSException as e:
            fetch_err = "{}".format(e)
        finally:
            result = {}
            result["error"] = fetch_err
            result["hostname"] = hostname
            result["record_type"] = record_type
            result["records"] = records

        return result


    def getARecords(self, hostname):
        return self.getRecordsBrief(hostname, "A")

    def getARecordsMT(self, hostname, mt_queue):
        mt_queue.put({"A": self.getRecordsBrief(hostname, "A")})

    def getCNAMERecords(self, hostname):
        return self.getRecordsBrief(hostname, "CNAME")

    def getCNAMERecordsMT(self, hostname, mt_queue):
        mt_queue.put({"CNAME": self.getRecordsBrief(hostname, "CNAME")})

    def getRecordsExhaustive(self, hostname, record_type):
        if record_type not in self.supported_records:
            raise InputError(record_type, "Unsupported record type: {}".format(record_type))

        depth = 0
        no_err = True
        allResults = []

        host = dns.name.from_text(hostname)

        while no_err and (depth < self.max_depth):
            try:
                records = []
                detailed_records = {}
                fetch_err = None
                cur_hostname = host.to_text()
                records = self.getRecords(cur_hostname, record_type)
                for record in records:
                    _record = record
                    if record_type == "MX":
                        _record = record.split()[1]
                    detailed_records[record] = self.getARecords(_record)
            except DNSException as e:
                fetch_err = "{}".format(e)
            finally:
                result = {}
                result["error"] = fetch_err
                result["hostname"] = cur_hostname
                result["record_type"] = record_type
                result["records"] = detailed_records
                allResults.append(result)

            try:
                host = host.parent()
                depth += 1
            except dns.name.NoParent:
                no_err = False

        return allResults


    def getMXRecords(self, hostname):
        return self.getRecordsExhaustive(hostname, "MX")

    def getMXRecordsMT(self, hostname, mt_queue):
        mt_queue.put({"MX": self.getMXRecords(hostname)})

    def getNSRecords(self, hostname):
        return self.getRecordsExhaustive(hostname, "NS")

    def getNSRecordsMT(self, hostname, mt_queue):
        mt_queue.put({"NS": self.getNSRecords(hostname)})

    def getAllRecords(self, hostname, records=["NS", "A", "MX", "CNAME"]):
        allRecords = {}
        if "A" in records:
            allRecords["A"] = self.getARecords(hostname)
        if "CNAME" in records:
            allRecords["CNAME"] = self.getCNAMERecords(hostname)
        if "MX" in records:
            allRecords["MX"] = self.getMXRecords(hostname)
        if "NS" in records:
            allRecords["NS"] = self.getNSRecords(hostname)
        return allRecords

    def getAllRecordsMT(self, hostname, records=["NS", "A", "MX", "CNAME"], nameserver=None):
        allRecords = {}
        threads = []
        mt_queue = Queue()
        if nameserver is not None:
            self.resolver.nameservers = nameserver
        if "A" in records:
            threads.append(Thread(target=self.getARecordsMT, args=(hostname, mt_queue)))
        if "CNAME" in records:
            threads.append(Thread(target=self.getCNAMERecordsMT, args=(hostname, mt_queue)))
        if "MX" in records:
            threads.append(Thread(target=self.getMXRecordsMT, args=(hostname, mt_queue)))
        if "NS" in records:
            threads.append(Thread(target=self.getNSRecordsMT, args=(hostname, mt_queue)))

        map(lambda thread: thread.start(), threads)
        map(lambda thread: thread.join(), threads)

        while not mt_queue.empty():
            allRecords.update(mt_queue.get())

        return allRecords