dnsdiff

#!/usr/bin/env python3

import argparse
import dns.query, dns.zone
import os.path
import socket
import sys

from blessings import Terminal
from random import randint
from time import sleep, strftime

# TODO: Package properly
VERSION = 'release-1.2'

def header(ns1_name, ns1_port, ns2_name, ns2_port):
    global seen_header

    std_port = socket.getservbyname('domain')
    t_stamp = strftime('%Y-%m-%d %H:%M:%S %z')

    # http://superuser.com/questions/710019/
    if not seen_header:
        if ns1_port != std_port or ns2_port != std_port:
            print("--- %s:%s\t%s\n+++ %s:%s\t\t%s" %
                  (ns1_name, ns1_port, t_stamp, ns2_name, ns2_port, t_stamp))
        else:
            print("--- %s\t%s\n+++ %s\t\t%s" %
                  (ns1_name, t_stamp, ns2_name, t_stamp))

    seen_header = True

def added(answer, color=False):
    for rrset in answer:
        if color:
            print(term.green + '+' + rrset + term.normal)
        else:
            print('+' + rrset)

def removed(answer, color=False):
    for rrset in answer:
        if color:
            print(term.red + '-' + rrset + term.normal)
        else:
            print('-' + rrset)

def get_response(query, ns_addr, ns_name, rname, rdtype, port):
    response = None
    successful = False
    t = 2

    while not successful:
        try:
            response = dns.query.udp(query, ns_addr, timeout=t, port=port)
            successful = True
        except dns.exception.Timeout:
            # Retry with backoff
            if t <= 8:
                t = t<<1
                continue
            else:
                break

    return response

def extract_rrset(response, ns_name, rdtype, ignore_ttl=False):
    # TODO: handle SERVFAIL, REFUSED, NOTIMP, etc?
    # Return only well-undestood cases so that this fails loudly

    # Successful query response
    if (response.rcode() == dns.rcode.NOERROR
            and response.answer):
        if ignore_ttl:
            response.answer[0].ttl = 0
        return sorted(str(response.answer[0]).split('\n'))

    # rrsets for NS records will be in the AUTHORITY section of the dns.message
    # object instead of the ANSWER section
    if (rdtype == 'NS'
            and response.rcode() == dns.rcode.NOERROR
            and not response.answer
            and not ns_name in response.authority[0]):
        if ignore_ttl:
            response.authority[0].ttl = 0
        return sorted(str(response.authority[0]).split('\n'))

    # AWS Alias record with NXDOMAIN target
    if (response.rcode() == dns.rcode.NOERROR
            and not response.answer):
        return None

def main():
    progname = os.path.basename(sys.argv[0])
    parser = argparse.ArgumentParser(prog=progname)
    parser._optionals.title = "Options"

    parser.add_argument("-V", "--version", action="version",
                        version="%(prog)s: " + VERSION)

    parser.add_argument("-c", "--color", dest="color_flag", action="store_true",
                        help="enable colorized output")

    parser.add_argument("-d", "--delay-max", dest="delay_max",
                        metavar='SECONDS',
                        help="maximum number of seconds of delay to introduce" +
                             " between each request")

    parser.add_argument("-t", "--ignore-ttl", dest="ttl_flag",
                        action="store_true",
                        help="ignore changes to TTL values")

    parser.add_argument("-f", "--zonefile", dest="filename", required=True,
                        help="FILENAME is expected to be a valid " +
                             "zonefile exported from NAMESERVER1\n" +
                             "https://tools.ietf.org/html/rfc1035#section-5")

    parser.add_argument("--from-ns", dest="ns1_name", required=True,
                        metavar="NAMESERVER1[:PORT]",
                        help="compare responses to NAMESERVER2")

    parser.add_argument("--to-ns", dest="ns2_name", required=True,
                        metavar="NAMESERVER2[:PORT]",
                        help="compare responses to NAMESERVER1")

    parser.add_argument("--origin", dest="origin",
                        help="Set origin of the zonefile")

    args = parser.parse_args()

    if args.color_flag:
        global term
        term = Terminal()

    try:
        zone = dns.zone.from_file(args.filename,
                                  origin=args.origin,
                                  allow_include=False,
                                  relativize=False)
    except dns.zone.UnknownOrigin:
        print("{}: Unable to find origin parameter."
              " Please use argument --origin.".format(progname))
        sys.exit(1)
    except dns.exception.DNSException:
        print("{}: Unable to import {}.".format(progname, args.filename))
        sys.exit(1)

    # Default port
    ns1_port = ns2_port = 53

    # Support for alternate port
    if ':' in args.ns1_name:
         ns1_port = int(args.ns1_name.split(':')[1])
         args.ns1_name = args.ns1_name.split(':')[0]

    if ':' in args.ns2_name:
         ns2_port = int(args.ns2_name.split(':')[1])
         args.ns2_name = args.ns2_name.split(':')[0]

    # TODO: IPv6 support with socket.getaddrinfo()
    ns1_addr = socket.gethostbyname(args.ns1_name)
    ns2_addr = socket.gethostbyname(args.ns2_name)

    global seen_header
    seen_header = False

    for (rname, rdataset) in zone.iterate_rdatasets():
        rdtype = dns.rdatatype.to_text(rdataset.rdtype)
        query = dns.message.make_query(rname, rdtype)

        # TODO: Improve accuracy for response pools. Execute each query
        # multiple times until most/all records are exposed and gathered.

        r1 = get_response(query, ns1_addr, args.ns1_name, rname, rdtype, ns1_port)

        if not r1:
            print("{}: Connection timed out to: '{}'".format(
                progname, args.ns1_name))
            print("Try using or increasing --delay-max")
            sys.exit(1)

        r2 = get_response(query, ns2_addr, args.ns2_name, rname, rdtype, ns2_port)

        if not r2:
            print("{}: Connection timed out to: '{}'".format(
                progname, args.ns2_name))
            print("Try using or increasing --delay-max")
            sys.exit(1)

        # FIXME: Please report any issues with answer extraction
        a1 = extract_rrset(r1, args.ns1_name, rdtype, ignore_ttl=args.ttl_flag)
        a2 = extract_rrset(r2, args.ns2_name, rdtype, ignore_ttl=args.ttl_flag)

        if a1 and a2:
            # Answers from both servers, compare rrsets
            if a1 != a2:
                header(args.ns1_name, ns1_port, args.ns2_name, ns2_port)
                removed(a1, color=args.color_flag)
                added(a2, color=args.color_flag)
        #   else:
        #       # Records both exist and match
        #       pass
        elif a1 and not a2:
            # Added to ns1 removed from ns2
            header(args.ns1_name, ns1_port, args.ns2_name, ns2_port)
            removed(a1, color=args.color_flag)
        elif a2 and not a1:
            # Added to ns2 removed from ns1
            header(args.ns1_name, ns1_port, args.ns2_name, ns2_port)
            added(a2, color=args.color_flag)
        else:
            # Exists in zonefile but not on either server
            print("{}: Record \"{}\" exists in zonefile but not on either"
                  "server.".format(progname, query.question[0]))
            print("Export zonefile from {} and try again".format(
                args.ns1_name))
            sys.exit(1)

        # Attempt to evade rate limits
        if args.delay_max:
            sleep(randint(0, args.delay_max))

if __name__ == '__main__':
    main()