ACN_Project/flow.py at master · TheodoreJBieber/ACN_Project · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
'''
Python class defining a flow.
A flow is a substructure within a packet burst.
A flow is defined by the source ip, destination ip, source port, destination port, and protocol on which the packets are being sent

A flow will contain information such as bytes sent, received,
'''
from time import strftime, localtime

from scapy.all import IP, TCP

class Flow:

    def __init__(self, source_ip=None, dest_ip=None, source_port=None, dest_port=None, protocol=None):
        self.init(source_ip, dest_ip, source_port, dest_port, protocol)

    def init(self, source_ip=None, dest_ip=None, source_port=None, dest_port=None, protocol=None):
        self.source_ip = source_ip
        self.dest_ip = dest_ip
        self.source_port = source_port
        self.dest_port = dest_port
        self.protocol = protocol

        # tracking
        self.timestamp = None
        self.packets_sent=0
        self.packets_received=0
        self.bytes_sent=0
        self.bytes_received=0

        self.protostr = ""


    '''
    Syntax: flowObject == packet
    returns: -true if the packet has the same source ip, dest ip, source port, dest port, and protocol as the flow
             -false otherwise
    '''
    def __eq__(self, packet): # TODO: This may change depending on whether the order matters for source/dest
        return (
            # check if A was source and B was dest
            ((self.source_ip == packet[IP].src and
            self.dest_ip == packet[IP].dst and
            self.source_port == packet[0].sport and
            self.dest_port == packet[0].dport and
            self.protocol == packet[IP].proto))
            or # vice versa
            ((self.dest_ip == packet[IP].src and
            self.source_ip == packet[IP].dst and
            self.dest_port == packet[0].sport and
            self.source_port == packet[0].dport and
            self.protocol == packet[IP].proto)))

    '''
    The inverse of __eq__
    Syntax: flowObject == packet
    returns: -false if the packet has the same source ip, dest ip, source port, dest port, and protocol as the flow
             -true otherwise
    '''
    def __ne__(self, packet):
        return not self == packet

    '''
    Syntax: flowObject + packet
    Adds a packet to the flow
    '''
    def __add__(self, packet): # TODO: still need to add sent/received stuff
        if self.source_ip == None:
            self.init(packet[IP].src,packet[IP].dst,packet[TCP].sport,packet[TCP].dport,packet[IP].proto)

            # from stack overflow
            try: # hope nothing goes wrong
                proto_field = IP().get_field('proto')
                self.protostr = proto_field.i2s[packet.proto]
            except:
                self.protostr = str(self.protocol) # backup is to use int


        flowstring = "<"+str(self.source_ip)+">"+"<"+str(self.dest_ip)+">"+"<"+str(self.source_port)+">"+"<"+str(self.dest_port)+">"+"<"+str(self.protocol)+">"
        packetstring = "<"+str(packet[IP].src)+">"+"<"+str(packet[IP].dst)+">"+"<"+str(packet[TCP].sport)+">"+"<"+str(packet[TCP].dport)+">"+"<"+str(packet[IP].proto)+">"
        assert (self == packet), "flow signature " + flowstring + " differs from packet signature " + packetstring

        # all good!
        # tracking
        self.timestamp = packet.time
        sent = self.source_ip == packet[IP].src

        if sent:
            self.packets_sent+=1
            self.bytes_sent+=len(packet)
        else:
            self.packets_received+=1
            self.bytes_received+=len(packet)

        return self # return the flow object for convenience

    '''
    Prints the current flow summary information
    '''
    def __str__(self):
        return (
            "<"+str(strftime('%Y-%m-%d %H:%M:%S', localtime(self.timestamp)))+">"+
            " <"+str(self.source_ip)+">"+
            " <"+str(self.dest_ip)+">"+
            " <"+str(self.source_port)+">"+
            " <"+str(self.dest_port)+">"+
            " <"+str(self.protostr)+">"+
            " <"+str(self.packets_sent)+">"+
            " <"+str(self.packets_received)+">"+
            " <"+str(self.bytes_sent)+">"+
            " <"+str(self.bytes_received)+">"
        )