import sys
import dpkt
import socket
import struct
import time
import os

def calcHost(hostname, type, srcIP, dstIP):
    chars = {}
    qs = []
    totLen = 0
    for ch in hostname:
        if not chars.has_key(ch):
            chars[ch] = 0
        chars[ch] += 1

    mehSrc = struct.unpack(">L", srcIP)[0]
    mehDst = struct.unpack(">L", dstIP)[0]
    outStr = "%f,%f" % (mehSrc/float(0xFFFFFFFF),mehDst/float(0xFFFFFFFF))

    outStr += ",%d,%d,%d,%d#%s;%s;%s" % (1, len(hostname), len(chars),
                                       type, socket.inet_ntoa(srcIP), socket.inet_ntoa(dstIP), hostname)
    s.send(outStr)
    sys.stdout.write(".")
    sys.stdout.flush()

def handlePacket(pktlen, data, ts):
    ether = dpkt.ethernet.Ethernet(data)
    if ether.type != dpkt.ethernet.ETH_TYPE_IP: return

    ip = ether.data
    if ip.p != dpkt.ip.IP_PROTO_UDP: return

    udp = ip.data
    if udp.dport != 53: return

    dns = dpkt.dns.DNS(udp.data)
    for q in dns.qd:
        try:
            calcHost(q.name, q.type, ip.src, ip.dst)
        except Exception, err:
            print "Err: %s" % err 

def readFromFile(fname):
    # change this to False to play back as quickly as possible
    realtime = True

    pcapReader = dpkt.pcap.Reader(file(fname, "rb"))
    start = time.time()
    firstPacketTS = None

    for ts, data in pcapReader:
        if realtime:
            if not firstPacketTS: firstPacketTS = ts
            while ts - firstPacketTS > time.time() - start:
                time.sleep(0.01)            
        handlePacket(len(data), data, ts)

def readFromWire(dev):
    import pcap
    pc = pcap.pcapObject()
    pc.open_live(dev, 1600, 0, 100)
    pc.setfilter("udp dst port 53", 0, 0)
    while 1:
        pc.dispatch(1, handlePacket)

if __name__ == "__main__":
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.connect(("127.0.0.1", 9000))

    if os.path.exists(sys.argv[1]):
        readFromFile(sys.argv[1])
    else:
        readFromWire(sys.argv[1])