diff --git a/src/scapy_tcp/tcp_state.py b/src/scapy_tcp/tcp_state.py index 7e72876..602990a 100644 --- a/src/scapy_tcp/tcp_state.py +++ b/src/scapy_tcp/tcp_state.py @@ -31,11 +31,14 @@ is_syn_pkt = lambda pkt: 'TCP' in pkt and pkt['TCP'].flags == TCP_FLAGS['S'] is_synack_pkt = lambda pkt: 'TCP' in pkt and pkt['TCP'].flags == (TCP_FLAGS['S'] | TCP_FLAGS['A']) -create_pkt_flow = lambda pkt: "%s:%s ==> %s:%s"%(pkt['IP'].src,str(pkt['IP'].sport),pkt['IP'].dst,str(pkt['IP'].dport)) +create_pkt_flow = lambda pkt: "%s:%s ==> %s:%s"%(pkt['IP'].src,str(pkt['IP'].sport),pkt['IP'].dst,str(pkt['IP'].dport)) \ + if 'IP' in pkt else "[%s]:%s ==> [%s]:%s"%(pkt['IPv6'].src,str(pkt['IPv6'].sport),pkt['IPv6'].dst,str(pkt['IPv6'].dport)) -create_forward_flow = lambda pkt: "%s:%s ==> %s:%s"%(pkt['IP'].src,str(pkt['IP'].sport),pkt['IP'].dst,str(pkt['IP'].dport)) +create_forward_flow = lambda pkt: "%s:%s ==> %s:%s"%(pkt['IP'].src,str(pkt['IP'].sport),pkt['IP'].dst,str(pkt['IP'].dport)) \ + if 'IP' in pkt else "[%s]:%s ==> [%s]:%s"%(pkt['IPv6'].src,str(pkt['IPv6'].sport),pkt['IPv6'].dst,str(pkt['IPv6'].dport)) -create_reverse_flow = lambda pkt: "%s:%s ==> %s:%s"%(pkt['IP'].dst,str(pkt['IP'].dport),pkt['IP'].src,str(pkt['IP'].sport)) +create_reverse_flow = lambda pkt: "%s:%s ==> %s:%s"%(pkt['IP'].dst,str(pkt['IP'].dport),pkt['IP'].src,str(pkt['IP'].sport)) \ + if 'IP' in pkt else "[%s]:%s ==> [%s]:%s"%(pkt['IPv6'].dst,str(pkt['IPv6'].dport),pkt['IPv6'].src,str(pkt['IPv6'].sport)) create_flow = create_forward_flow @@ -73,9 +76,13 @@ def init(self, pkt): raise Exception("Not valid SYN") self.flows = set((create_forward_flow(pkt), create_reverse_flow(pkt))) - self.server = pkt['IP'].dst - self.client = pkt['IP'].src - + if 'IP' in pkt: + self.server = pkt['IP'].dst + self.client = pkt['IP'].src + else: + self.server = pkt['IPv6'].dst + self.client = pkt['IPv6'].src + # 0 is now, 1 is the future Flags self.server_state = "LISTEN" self.client_state = "SYN_SENT" @@ -95,7 +102,7 @@ def next_state(self, pkt): if flow not in self.flows: raise Exception("Not a valid packet for this model") - if pkt['IP'].dst == self.server: + if pkt['IP' if 'IP' in pkt else 'IPv6'].dst == self.server: v = self.handle_client_pkt(pkt) if self.is_fin_wait(): self.fin_wait_time = pkt.time @@ -357,7 +364,10 @@ def build_basic_pkt(self, dst, dport, sport=None): self.sport = sport self.dport = dport self.dst = dst - return IP(dst=dst) / TCP(dport=dport, sport=sport) + if ':' in dst: + return IPv6(dst=dst) / TCP(dport=dport, sport=sport) + else: + return IP(dst=dst) / TCP(dport=dport, sport=sport) def get_rbase_tcp(self, rseg): ''' @@ -381,7 +391,10 @@ def get_rbase_ip(self, rpkt): dst = rpkt.src src = rpkt.dst options = rpkt.options - return IP(src=src, dst=dst, options=options) + if ':' in src: + return IPv6(src=src, dst=dst, options=options) + else: + return IP(src=src, dst=dst, options=options) def get_rbase_pkt(self, rpkt): ''' @@ -390,7 +403,10 @@ def get_rbase_pkt(self, rpkt): :param rpkt: rcvd segment to base a new packet off of :type rpkt: scapy.IP/scapy.TCP ''' - return IP(dst=rpkt[IP].src) / TCP(dport=rpkt[TCP].sport, sport=rpkt[TCP].dport) + if IP in rpkt: + return IP(dst=rpkt[IP].src) / TCP(dport=rpkt[TCP].sport, sport=rpkt[TCP].dport) + else: + return IPv6(dst=rpkt[IPv6].src) / TCP(dport=rpkt[TCP].sport, sport=rpkt[TCP].dport) def get_base_tcp(self): ''' @@ -406,13 +422,19 @@ def get_base_ip(self): Creates a base IP packet based on internal TCP/IP stuffs. ''' dst = self.dst - return IP(dst=dst) + if ':' in dst: + return IPv6(dst=dst) + else: + return IP(dst=dst) def get_base_pkt(self): ''' Creates a base packet based on a rcvd packet. ''' - return IP(dst=self.dst) / TCP(dport=self.dport,sport=self.sport) + if ':' in self.dst: + return IPv6(dst=self.dst) / TCP(dport=self.dport,sport=self.sport) + else: + return IP(dst=self.dst) / TCP(dport=self.dport,sport=self.sport) def update_seg_state(self, seg, payload=None): diff --git a/src/scapy_tcp/tcp_stream.py b/src/scapy_tcp/tcp_stream.py index eae2192..7145172 100644 --- a/src/scapy_tcp/tcp_stream.py +++ b/src/scapy_tcp/tcp_stream.py @@ -43,8 +43,12 @@ class TCPStream: def __init__(self, pkt ): - self.src = pkt["IP"].src - self.dst = pkt["IP"].dst + if "IP" in pkt: + self.src = pkt["IP"].src + self.dst = pkt["IP"].dst + else: + self.src = pkt["IPv6"].src + self.dst = pkt["IPv6"].dst self.sport = pkt["TCP"].sport self.dport = pkt["TCP"].dport self.time = float(pkt.time) @@ -92,16 +96,28 @@ def add_pkt(self, pkt): return is_closed def create_client_directed_flow(self): - return "%s:%s ==> %s:%s"%(self.src,str(self.sport),self.dst,str(self.dport)) + if ':' in self.src: + return "[%s]:%s ==> [%s]:%s"%(self.src,str(self.sport),self.dst,str(self.dport)) + else: + 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),) + if ':' in self.dst: + return "[%s]:%s ==> [%s]:%s"%(self.dst,str(self.dport),self.src,str(self.sport),) + else: + 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)) + if ':' in self.src: + return "[%s]:%s ==> [%s]:%s"%(self.src,str(self.sport),self.dst,str(self.dport)) + else: + 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)) + if ':' in self.dst: + return "[%s]:%s <== [%s]:%s"%(self.dst,str(self.dport),self.src,str(self.sport)) + else: + return "%s:%s <== %s:%s"%(self.dst,str(self.dport),self.src,str(self.sport)) def get_client_server(self): return 0 @@ -147,7 +163,7 @@ def get_app_stream_summary(self, pkts_cnt=0): flow_total += payload_len time_elapsed, time_last_pkt = self.packet_time_spacing_idx(pkts, i) - if self.src == pkt['IP'].src: + if self.src == pkt['IP' if 'IP' in pkt else 'IPv6'].src: #flow_info = self.get_server_client_str() flow_info = self.get_client_server() if not last_client_pkt is None: @@ -195,7 +211,7 @@ def get_stream_summary(self, pkts_cnt=0): flow_total += payload_len time_elapsed, time_last_pkt = self.packet_time_spacing_idx(pkts, i) - if self.src == pkt['IP'].src: + if self.src == pkt['IP' if 'IP' in pkt else 'IPv6'].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) @@ -239,22 +255,22 @@ def get_ip_summary(self, pkts_cnt=0): while i < len(pkts): pkt = pkts[i] flow_info = '' - payload_len = len(pkt['IP']) + payload_len = len(pkt['IP' if 'IP' in pkt else 'IPv6']) flow_total += payload_len time_elapsed, time_last_pkt = self.packet_time_spacing_idx(pkts, i) - if self.src == pkt['IP'].src: + if self.src == pkt['IP' if 'IP' in pkt else 'IPv6'].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']) + client_total += len(pkt['IP' if 'IP' in pkt else 'IPv6']) 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']) + server_total += len(pkt['IP' if 'IP' in pkt else 'IPv6']) pkt_summary.append([ str(pkt.time),