Support splitting of midstream TCP flows.

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

3 files changed, 14 insertions, 9 deletions

diff --git a/TCPSplit.py b/TCPSplit.py
index e85a41f..391d0a9 100755
--- a/TCPSplit.py
+++ b/TCPSplit.py
@@ -39,12 +39,18 @@ class TCPSplitStream(object):
                             ', expected ' + str(TCPStream.TCPStream))
         self.stream = tcp_stream
         self.ordered_pkts = self.stream.get_order_pkts()
-
-    def __generate_handshake(self):
-        self.seq = self.ordered_pkts[0][scapy.all.TCP].seq # TCP-SYN
-        self.ack = self.ordered_pkts[1][scapy.all.TCP].seq # TCP-SYN-ACK
         self.ip2dst = scapy.all.IP(src = self.stream.src, dst = self.stream.dst)
         self.ip2src = scapy.all.IP(src = self.stream.dst, dst = self.stream.src)
+        if self.stream.tcp_state.syn_seen is True:
+            self.seq = self.ordered_pkts[0][scapy.all.TCP].seq # TCP-SYN
+            self.ack = self.ordered_pkts[1][scapy.all.TCP].seq # TCP-SYN-ACK
+        else:
+            self.seq = self.ordered_pkts[0][scapy.all.TCP].seq
+            self.ack = self.ordered_pkts[1][scapy.all.TCP].ack
+
+    def __generate_handshake(self):
+        if self.stream.tcp_state.syn_seen is False:
+            return list()
 
         syn = scapy.all.TCP(sport = self.stream.sport, dport = self.stream.dport,
                             flags = 'S', seq = self.seq, ack = 0)
diff --git a/TCPState.py b/TCPState.py
index 2c936d9..9a36eab 100644
--- a/TCPState.py
+++ b/TCPState.py
@@ -67,9 +67,8 @@ class TCPStateMachine:
     def init(self, pkt):
         if not 'TCP' in pkt:
             raise Exception("Not a TCP Packet")
-        if not is_syn_pkt(pkt):
-            raise Exception("Not valid SYN")
-            
+
+        self.syn_seen = is_syn_pkt(pkt)
         self.flows = set((create_forward_flow(pkt), create_reverse_flow(pkt)))
         self.server = pkt['IP'].dst
         self.client = pkt['IP'].src
diff --git a/TCPStreamExtractor.py b/TCPStreamExtractor.py
index 91b7fa9..d6dc10a 100644
--- a/TCPStreamExtractor.py
+++ b/TCPStreamExtractor.py
@@ -104,13 +104,13 @@ class TCPStreamExtractor:
 
         flow = (create_forward_flow(pkt), create_reverse_flow(pkt))
         if not flow[0] in self.streams and\
-            not flow[1] in self.streams and is_syn_pkt(pkt):
+            not flow[1] in self.streams:
             self.streams [flow[0]] = TCPStream(pkt)
             self.streams [flow[1]] = self.streams [flow[0]]
             self.fwd_flows.add(flow[0])
             self.rev_flows.add(flow[1])
         elif flow[0] in self.streams:
-            self.streams[flow[0]].add_pkt(pkt)        
+            self.streams[flow[0]].add_pkt(pkt)
         return pkt
 
     def process_packets(self):