Initial commit.

Signed-off-by: lns <matzeton@googlemail.com>

1 files changed, 332 insertions, 0 deletions

diff --git a/TCPStream.py b/TCPStream.py
new file mode 100644
index 0000000..e460b6c
--- /dev/null
+++ b/TCPStream.py
@@ -0,0 +1,332 @@
+"""
+TCP stream extraction using Scapy.
+
+(c) Praetorian 
+Author: Adam Pridgen <adam.pridgen@praetorian.com> || <adam.pridgen@thecoverofnight.com>
+
+
+This program is free software; you can redistribute it and/or modify it
+under the terms of the GNU General Public License as published by the Free
+Software Foundation; either version 3, or (at your option) any later
+version.
+
+This program is distributed in the hope that it will be useful, but WITHOUT
+ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+more details.
+
+You should have received a copy of the GNU General Public License along
+with this program; see the file COPYING.  If not, write to the Free
+Software Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
+02110-1301, USA.
+
+description: keeps track of the packets belonging to a specific tcp stream
+    
+"""
+
+from scapy.all import *
+from random import randint
+import scapy
+
+from random import randint
+from TCPState import TCPStateMachine
+
+
+from threading import *
+
+HEADER = "Time Flow_with_Direction PktPayload ClientDataTransfered ServerDataTransfered TotalData TimeBetweenLastPkt TimeBetweenLastClientPacket TimeBetweenLastServerPacket  TotalTimeElapsed"
+
+
+MAX = 2 ** 32 - 1
+
+FLOW_NAME = "T:%s_C:%s:%s_S:%s:%s"
+
+class TCPStream:
+    def __init__(self, pkt ):
+        self.src = pkt["IP"].src 
+        self.dst = pkt["IP"].dst
+        self.sport = pkt["TCP"].sport
+        self.dport = pkt["TCP"].dport        
+        self.time = float(pkt.time)
+        self.tcp_state = TCPStateMachine(pkt)
+        self.flows = self.tcp_state.flows
+        self.pkts = {self.time:pkt}
+        self.ordered_pkts  = None
+        self.last_time_written = 0.0
+        self.pcap_file = None
+        self.couchdb_file = None
+        self.flow_file = None
+
+    def len_pkts(self):
+        return len(list(self.pkts.keys()))
+    
+    def destroy(self, pkts_cnt=0):
+        #self.ordered_pkts = None       
+        keys = list(self.pkts.keys())
+        print(("destroying packets for %s"%self.get_stream_name()))
+        if pkts_cnt > 0:
+           keys = keys[:pkts_cnt]
+           print(("\t...Destroying the first %d pkts"%pkts_cnt))
+        
+        if not self.ordered_pkts is None and len(self.ordered_pkts) >= pkts_cnt:
+            for i in self.ordered_pkts[:pkts_cnt]:
+                del i
+        
+        self.ordered_pkts = None
+        
+        for key in keys:
+            del self.pkts[key]
+
+    def get_stream_name(self):
+        global FLOW_NAME
+        return FLOW_NAME%(self.time, self.src, self.sport, self.dst, self.dport)
+        
+    def get_stream_keys(self):
+        global HEADER
+        return HEADER.split()
+    
+    def add_pkt(self, pkt):
+        is_closed = self.tcp_state.next_state(pkt)
+        self.pkts[float(pkt.time)] = pkt
+        self.last_time = pkt.time
+        return is_closed
+        
+    def create_client_directed_flow(self):
+        return "%s:%s ==> %s:%s"%(self.src,str(self.sport),self.dst,str(self.dport))
+    
+    def create_server_directed_flow(self):
+        return "%s:%s ==> %s:%s"%(self.dst,str(self.dport),self.src,str(self.sport),)
+
+    def get_client_server_str(self):
+        return "%s:%s ==> %s:%s"%(self.src,str(self.sport),self.dst,str(self.dport))
+
+    def get_server_client_str(self):
+        return "%s:%s <== %s:%s"%(self.dst,str(self.dport),self.src,str(self.sport))
+
+    def get_client_server(self):
+        return 0
+
+    def get_server_client(self):
+        return 1
+    
+    def interp_direction(self, dir):
+        if dir == 0:
+           return get_client_server_str()
+        elif dir == 1:
+           return get_server_client_str()
+        raise Exception("Not a valid direction for the client server")
+
+    def get_order_pkts(self, pkts_cnt=0):
+        # caching mechanism
+        if self.ordered_pkts is None or len(self.ordered_pkts) != len(list(self.pkts.keys())):
+            times = list(self.pkts.keys())
+            times.sort()
+            self.ordered_pkts = [self.pkts[time] for time in times]
+        if pkts_cnt > 0:
+           return self.ordered_pkts[:pkts_cnt]
+        return self.ordered_pkts
+
+    def get_app_stream_summary(self, pkts_cnt=0):
+        global HEADER
+        pkts = self.get_order_pkts(pkts_cnt)
+        pkt_summary = []#[HEADER]
+        flow_total = 0
+        client_total = 0.0
+        server_total = 0.0
+        
+        last_client_pkt =  None
+        last_server_pkt = None
+        last_client_pkt_time = 0.0
+        last_server_pkt_time = 0.0
+        
+        i = 0
+        while i < len(pkts):
+            pkt = pkts[i]
+            flow_info = ''
+            payload_len = len(pkt['TCP'].payload)
+            flow_total += payload_len
+            time_elapsed, time_last_pkt = self.packet_time_spacing_idx(pkts, i)
+            
+            if self.src == pkt['IP'].src:
+                #flow_info = self.get_server_client_str()
+                flow_info = self.get_client_server()
+                if not last_client_pkt is None:
+                    last_client_pkt_time = self.packet_time_spacing_pkt(pkt, last_client_pkt)
+                last_client_pkt = pkt
+                client_total += len(pkt['TCP'].payload)
+            else:
+                #flow_info = self.get_client_server_str()
+                flow_info = self.get_server_client()
+                if not last_server_pkt is None:
+                    last_server_pkt_time = self.packet_time_spacing_pkt(pkt,last_server_pkt)
+                last_server_pkt = pkt
+                server_total += len(pkt['TCP'].payload)
+
+            pkt_summary.append([ str(pkt.time), 
+                                        flow_info, 
+                                        str(payload_len), 
+                                        str(client_total),
+                                        str(server_total),
+                                        str(flow_total),
+                                        str(time_last_pkt),
+                                        str(last_client_pkt_time),
+                                        str(last_server_pkt_time),
+                                        str(time_elapsed)])
+            i+=1
+        return pkt_summary
+            
+    def get_stream_summary(self, pkts_cnt=0):
+        pkts = self.get_order_pkts(pkts_cnt)
+        pkt_summary = []# [HEADER]
+        flow_total = 0
+        client_total = 0.0
+        server_total = 0.0
+        
+        last_client_pkt =  None
+        last_server_pkt = None
+        last_client_pkt_time = 0.0
+        last_server_pkt_time = 0.0
+        
+        i = 0
+        while i < len(pkts):
+            pkt = pkts[i]
+            flow_info = ''
+            payload_len = len(pkt['TCP'])
+            flow_total += payload_len
+            time_elapsed, time_last_pkt = self.packet_time_spacing_idx(pkts, i)
+            
+            if self.src == pkt['IP'].src:
+                flow_info = self.get_server_client_str()
+                if not last_client_pkt is None:
+                    last_client_pkt_time = self.packet_time_spacing_pkt(pkt, last_client_pkt)
+                last_client_pkt = pkt
+                client_total += len(pkt['TCP'])
+            else:
+                flow_info = self.get_client_server_str()
+                if not last_server_pkt is None:
+                    last_server_pkt_time = self.packet_time_spacing_pkt(pkt,last_server_pkt)
+                last_server_pkt = pkt
+                server_total += len(pkt['TCP'])
+
+
+            pkt_summary.append([ str(pkt.time), 
+                                        flow_info, 
+                                        str(payload_len), 
+                                        str(client_total),
+                                        str(server_total),
+                                        str(flow_total),
+                                        str(time_last_pkt),
+                                        str(last_client_pkt_time),
+                                        str(last_server_pkt_time),
+                                        str(time_elapsed)])
+            i+=1
+        return pkt_summary
+
+            
+    def get_ip_summary(self, pkts_cnt=0):
+        pkts = self.get_order_pkts(pkts_cnt)
+        pkt_summary = []# [HEADER]
+        flow_total = 0
+        client_total = 0
+        server_total = 0
+        
+        last_client_pkt =  None
+        last_server_pkt = None
+        last_client_pkt_time = 0.0
+        last_server_pkt_time = 0.0
+        
+        i = 0
+        while i < len(pkts):
+            pkt = pkts[i]
+            flow_info = ''
+            payload_len = len(pkt['IP'])
+            flow_total += payload_len
+            time_elapsed, time_last_pkt = self.packet_time_spacing_idx(pkts, i)
+            
+            if self.src == pkt['IP'].src:
+                flow_info = self.get_server_client_str()
+                if not last_client_pkt is None:
+                    last_client_pkt_time = self.packet_time_spacing_pkt(pkt, last_client_pkt)
+                last_client_pkt = pkt
+                client_total += len(pkt['IP'])
+            else:
+                flow_info = self.get_client_server_str()
+                if not last_server_pkt is None:
+                    last_server_pkt_time = self.packet_time_spacing_pkt(pkt,last_server_pkt)
+                last_server_pkt = pkt
+                server_total += len(pkt['IP'])
+
+
+            pkt_summary.append([ str(pkt.time), 
+                                        flow_info, 
+                                        str(payload_len), 
+                                        str(client_total),
+                                        str(server_total),
+                                        str(flow_total),
+                                        str(time_last_pkt),
+                                        str(last_client_pkt_time),
+                                        str(last_server_pkt_time),
+                                        str(time_elapsed)])
+            i+=1
+        return pkt_summary
+
+
+    def packet_time_spacing_idx(self, pkts, idx):
+        if idx < 1:      return 0.0, 0.0
+        return self.packet_time_spacing_pkt(pkts[idx], pkts[0]), self.packet_time_spacing_pkt(pkts[idx], pkts[idx-1])
+    
+    def packet_time_spacing_pkt(self, pkt_a, pkt_b):
+        return pkt_a.time - pkt_b.time
+    
+
+    def write_flow(self, filename, pkts_cnt=0):
+        if self.flow_file is None:
+           self.flow_file =open(filename+".app_flow", 'w') 
+        
+        flow_infos = []
+        flows = self.get_app_stream_summary(pkts_cnt)
+        for flow_info in flows:
+            flow_infos.append(" ".join([str(i) for i in flow_info]))
+            
+        self.flow_file.write("\n".join(flow_infos)+"\n")
+        self.flow_file.flush()
+    
+    def write_pcap(self, filename, pkts_cnt=0):
+        if self.pcap_file is None:
+           self.pcap_file = scapy.utils.PcapWriter(filename+".pcap")
+        
+        pkts = self.get_order_pkts(pkts_cnt)
+        self.pcap_file.write(pkts)
+        self.pcap_file.flush()
+        
+        
+    def organize_streams(self):
+        data = {}
+        seq_nos = {}
+        for pkt in self.get_order_pkts():
+            tcpp = pkt[TCP]
+            seq = tcpp.seq
+            sport = tcpp.sport
+            if not sport in data:
+                data[sport] = []
+                seq_nos[sport] = []
+            
+            if seq in seq_nos[sport]:
+                # retransmit
+                continue
+
+            if Raw in tcpp:
+                data[sport].append([tcpp.seq, tcpp[Raw].load])
+                seq_nos[sport].append(seq)
+        return seq_nos, data
+
+    def get_stream_data(self):
+        seq_nos, data = self.organize_streams()
+        streams = {k:bytes() for k in data.keys()}
+
+        for sport, seqs_payloads in data.items():
+            stream = b''
+            for seq, payload in seqs_payloads:
+                stream = stream + payload
+            streams[sport] = stream
+        return streams