#! /usr/bin/python

import sys, socket, getopt

import dns.zone
import dns.query
import dns.message
import dns.rdatatype
import dns.rdataclass

def min_size(message):
    """Calculate the minimum size this response can be without setting
    the TC (truncation bit).  This is calculated by removing all of
    the additional section save the OPT pseudo-RR."""
    
    additional = message.additional
    # note that in dns.message.Message, the OPT RR is not stored in
    # the additional section (where it was).  Instead it is parsed
    # into various member variables, so we don't need to do anything
    # special to preserve it.
    message.additional = []
    wire = message.to_wire()
    message.additional = additional
    return len(wire)

def additional_qname_length(qname):
    """Calculate the additional size a message might take if it had
    the maximum sized qname (e.g., a 255-octet qname).  Since the
    message already contains some qname, calculate the difference."""
    return 255 - len(qname.to_wire())

def do_query(qname, qtype):
    """Given a qname and a qtype, use the dedicated socket to query a
    namesrver with DO=1 and BUFSIZE=4096, then calculate the full
    response size, the minimum size, and the additional qname
    delta."""

    q = dns.message.make_query(qname, qtype)
    q.want_dnssec(True)
    q.payload = 4096
    wire = q.to_wire()

    # send the query
    s.send(wire)
    # wait for the response (no timeouts yet)
    resp_wire = s.recv(65535)
    # calculate the base length.
    resp_len = len(resp_wire)
    # parse the response
    resp = dns.message.from_wire(resp_wire)
    # calculate our minimum size
    minimum = min_size(resp)
    # calculate our maximum size
    plus_qname = additional_qname_length(q.question[0].name)
    return (resp, resp_len, minimum, plus_qname)

class Histogram:
    """A class representing a basic histogram for DNS response
    sizes."""

    def __init__(self, factor):
        self.hist_ = {}
        self.labels_ = {}
        self.factor_ = factor
        self.minimum_ = 0
        self.maximum_ = 0
        self.minlabel_ = ""
        self.maxlabel_ = ""

    def add(self, length, label):
        if label in self.labels_: return
        self.labels_[label] = 1

        index = int(length / self.factor_)
        if index not in self.hist_:
            self.hist_[index] = 0

        self.hist_[index] += 1

        if self.minimum_ == 0 or length < self.minimum_:
            self.minimum_ = length
            self.minlabel_ = label
        if length > self.maximum_:
            self.maximum_ = length
            self.maxlabel_ = label

    def to_text(self):
        res = "range [%d - %d] min: %s, max: %s\n" % \
            (self.minimum_, self.maximum_, self.minlabel_, self.maxlabel_)
        keys = self.hist_.keys()
        keys.sort()
        for index in keys:
            rangestr = "[%d - %d)" % (index * self.factor_, 
                                      (index + 1) * self.factor_)
            res += "%20s : %d\n" % (rangestr, self.hist_[index])
        return res

def process_zone(zone, skip_addr=True):
    last_name = ""; last_type = 0

    for name, rd in zone.iterate_rdatasets():
        n = name.to_text()
        t = rd.rdtype
        
        # skip queries to the glue records; they are redundant with
        # the NS queries.
        if skip_addr and (t == dns.rdatatype.A or t == dns.rdatatype.AAAA):
            continue
        label = n + "/" + dns.rdatatype.to_text(t)
 
        (response, resp_len, minimum, plus_qname) = do_query(n, t)
        print "%d\t%d\t%d\t%s" % (resp_len, minimum, plus_qname, label)

        # Populate the referral histograms
        if t == dns.rdatatype.NS:
            ref_hist.add(minimum + plus_qname, label)
            max_ref_hist.add(resp_len + plus_qname, label)
        else:
            other_hist.add(resp_len, label)

        if n != last_name:
            label = n + "/ANY"
            (response, resp_len, minimum, plus_qname) = \
                do_query(n , dns.rdatatype.ANY)
            print "%d\t%d\t%d\t%s" % (resp_len, minimum, plus_qname, label)
            other_hist.add(resp_len, label)

            label = n + "_/A"
            (response, resp_len, minimum, plus_qname) = \
                do_query(n + "_", dns.rdatatype.A)
            print "%d\t%d\t%d\t%s" % (resp_len, minimum, plus_qname, label)
            # populate the nxdomain histogram
            nx_hist.add(resp_len + plus_qname, label)

            last_name = n

if __name__ == "__main__":

    port = 53
    host = "127.0.0.1"

    options, args = getopt.getopt(sys.argv[1:], "p:s:")
    for k,v in options:
        if k == '-p':
            port = int(v)
        elif k == '-s':
            host = v
            
    # create a query socket
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    # connect it to our server
    s.connect((host, port))

    ref_hist = Histogram(64)
    max_ref_hist = Histogram(64)
    nx_hist = Histogram(64)
    other_hist = Histogram(64)

    for zonefile in args:
        origin=None
        skip = True
        if ':' in zonefile:
            zonefile, origin = zonefile.split(':', 1)
        if origin and ':' in origin:
            origin, skip = origin.split(':', 1)
            skip = len(skip) == 0
        if origin and origin == '.': 
            rel = True
        else: 
            rel = False

        zone = dns.zone.from_file(zonefile, origin=origin, relativize=rel)
        
        process_zone(zone, skip)

    s.close()

    print >> sys.stderr, "Referral sizes (Maximum truncation sizes):"
    print >> sys.stderr, ref_hist.to_text()

    print >> sys.stderr, "Referral sizes (Maximum overall sizes):"
    print >> sys.stderr, max_ref_hist.to_text()

    print >> sys.stderr, "NXDOMAIN sizes (Maximum truncation sizes):"
    print >> sys.stderr, nx_hist.to_text()

    print >> sys.stderr, "Other response sizes (Full response sizes):"
    print >> sys.stderr, other_hist.to_text()