detector.py

##############################
# Name: Kyle Holzinger
# Date: April 20th, 2015
# CS 558 Lab 4
##############################

import dpkt, sys, socket

class Detector(object):

    def __init__(self, name=None):
        if len(sys.argv) == 2:
            self.fname = sys.argv[1]
        else:
            self.fname = name
        if self.fname is None:
            raise Exception('Usage: python detector.py test.pcap')
        self.file = open(self.fname, 'rb')
        self.pcap = dpkt.pcap.Reader(self.file)

        self.counts = {}

    def getIPs(self):
        print ('Counting packets...')
        counter = 0
        for ts, buf in self.pcap:
            counter += 1
            if counter % 100000 == 0:
                print('Counted %d packets'%counter)
            try:
                eth = dpkt.ethernet.Ethernet(buf)
                ip = eth.data
                tcp = ip.data
                src = ip.src
                dst = ip.dst
                # if flags == 2 then SYN
                if tcp.flags == 2:
                    if self.counts.get(src) is None:
                        self.counts[src] = [0, 0]
                    self.counts[src][0] += 1

                # if flags == 18 then SYN, ACK
                elif tcp.flags == 18:
                    if self.counts.get(dst) is None:
                        self.counts[dst] = [0, 0]
                    self.counts[dst][1] += 1

            except Exception:
                continue
        print ('Finished counting')
        for key in self.counts.keys():
            # Edit this to remove the ones that were not valid.
            if not(self.counts[key][0] >= 3*self.counts[key][1]) or self.counts[key][0] == 1:
                del self.counts[key]

        return [socket.inet_ntoa(x) for x in self.counts.keys()]

'''
test.pcap is the test file that was provided with the assignment. however, this
function will prefer to use the filename passed in to the python script.
'''

if __name__ == '__main__':
    part3 = Detector()
    print(part3.getIPs())