#!/usr/bin/env python

import dns.message
import dns.rrset
import dns.flags
import dns.name
import btk
import socket
import sys
import struct
from optparse import OptionParser


__author__ = "Nominum, Inc."
__version__ = "1.0.0.0"
__date__ = "2006-09-28" 

IPHeader = '!BBHHHBBHLL'

IPHDRSTART = 14
IPHDRLEN = 20
UDPHDRLEN = 8
qtypecount = {}

def main(argv):
    parser = OptionParser(usage="%prog [-i] [-o]",
                          version = "%prog " + __version__ )
    parser.add_option("-i", "--input", dest="fin",
                     help="name of tcpdump file to parse", metavar="FILE")
    parser.add_option("-o", "--output", dest="fout",
                     help="file in which to save parsed DNS queries",
                     metavar="FILE")
    (opts, args) = parser.parse_args()
    if not opts.fin or not opts.fout:
        parser.error("incorrect number of arguments")

    pcap = btk.pcap()
    pcap.open_offline(opts.fin)
    outfile = open(opts.fout, "w")
    while True:
        packet = pcap.next()
        if packet is None:
            break
    
        # Toss the stuff before the IP header
        packet = packet[IPHDRSTART:]

        # Grab the rest of the packet so we can parse proto 
        iphdr = packet[0:IPHDRLEN]
        (vhl, tos, tlen, ipid, fragoff, ttl, proto, cksum, srcip, dstip) = \
                struct.unpack(IPHeader, iphdr)
 
        # Toss the IP header, we're done with it.  We need to account
        # for any IP header options.
        ihl = (vhl & 0xF) * 4
        packet = packet[ihl:]

        if proto == socket.IPPROTO_UDP: # UDP, 8-byte header
            packet = packet[UDPHDRLEN:]
        else:
            continue

        try:
            msg = dns.message.from_wire(packet)
        except:
            continue
        if not (msg.flags & dns.flags.QR):
            for query in msg.question: # handle multiple queries per packet
                fqdn = query.name.to_text()
                qtype = dns.rdatatype.to_text(query.rdtype)
                outfile.write("%s %s\n" % (fqdn, qtype))
                # add qtype to dict if not present, otherwise increment 
                qtypecount[qtype] = qtypecount.get(qtype, 0) + 1 

    outfile.close()
    sum = 0
    print "Statistics:"
    for v, d in qtypecount.items():
        if d:
            print "    %10s:\t%d" % (v, d)
            sum += d
    print "-------------------------"
    print "         TOTAL:\t%d" % sum

if __name__ == '__main__':
    main(sys.argv[1:])

