python: tcp.py: add state2str() and apply to tools

There are two tools/{tcpretrans.py,tcpstates.py,tcpdrop.py} display the
TCP state, and those code already exist in bcc/tcp.py, thus, this patch
remove the duplicate code and add bcc.tcp.state2str() function.

Signed-off-by: Rong Tao <[email protected]>

Author: Rtoax

src/python/bcc/tcp.py

@@ -13,19 +13,19 @@
 # limitations under the License.
 
 # from include/net/tcp_states.h:
-tcpstate = {}
-tcpstate[1] = 'ESTABLISHED'
-tcpstate[2] = 'SYN_SENT'
-tcpstate[3] = 'SYN_RECV'
-tcpstate[4] = 'FIN_WAIT1'
-tcpstate[5] = 'FIN_WAIT2'
-tcpstate[6] = 'TIME_WAIT'
-tcpstate[7] = 'CLOSE'
-tcpstate[8] = 'CLOSE_WAIT'
-tcpstate[9] = 'LAST_ACK'
-tcpstate[10] = 'LISTEN'
-tcpstate[11] = 'CLOSING'
-tcpstate[12] = 'NEW_SYN_RECV'
+_tcpstate = {}
+_tcpstate[1] = 'ESTABLISHED'
+_tcpstate[2] = 'SYN_SENT'
+_tcpstate[3] = 'SYN_RECV'
+_tcpstate[4] = 'FIN_WAIT1'
+_tcpstate[5] = 'FIN_WAIT2'
+_tcpstate[6] = 'TIME_WAIT'
+_tcpstate[7] = 'CLOSE'
+_tcpstate[8] = 'CLOSE_WAIT'
+_tcpstate[9] = 'LAST_ACK'
+_tcpstate[10] = 'LISTEN'
+_tcpstate[11] = 'CLOSING'
+_tcpstate[12] = 'NEW_SYN_RECV'
 
 # from include/net/tcp.h:
 TCPHDR_FIN = 0x01
@@ -56,3 +56,6 @@ def flags2str(flags):
     if flags & TCPHDR_CWR:
         arr.append("CWR")
     return "|".join(arr)
+
+def state2str(state):
+    return _tcpstate.get(state, str(state))

tools/tcpdrop.py

@@ -342,7 +342,7 @@
 def print_ipv4_event(cpu, data, size):
     event = b["ipv4_events"].event(data)
     reason_str = drop_reasons.get(event.drop_reason, "UNKNOWN")
-    state_flag_str = "%s (%s)" % (tcp.tcpstate[event.state], tcp.flags2str(event.tcpflags))
+    state_flag_str = "%s (%s)" % (tcp.state2str(event.state), tcp.flags2str(event.tcpflags))
     print("%-8s %-7d %-2d %-20s > %-20s %-20s %s (%d)" % (
         strftime("%H:%M:%S"), event.pid, event.ip,
         "%s:%d" % (inet_ntop(AF_INET, pack('I', event.saddr)), event.sport),
@@ -356,7 +356,7 @@ def print_ipv4_event(cpu, data, size):
 def print_ipv6_event(cpu, data, size):
     event = b["ipv6_events"].event(data)
     reason_str = drop_reasons.get(event.drop_reason, "UNKNOWN")
-    state_flag_str = "%s (%s)" % (tcp.tcpstate[event.state], tcp.flags2str(event.tcpflags))
+    state_flag_str = "%s (%s)" % (tcp.state2str(event.state), tcp.flags2str(event.tcpflags))
     print("%-8s %-7d %-2d %-20s > %-20s %-20s %s (%d)" % (
         strftime("%H:%M:%S"), event.pid, event.ip,
         "%s:%d" % (inet_ntop(AF_INET6, event.saddr), event.sport),

tools/tcpretrans.py

@@ -16,7 +16,7 @@
 # 03-Nov-2017   Matthias Tafelmeier Extended this.
 
 from __future__ import print_function
-from bcc import BPF
+from bcc import BPF, tcp
 import argparse
 from time import strftime
 from socket import inet_ntop, AF_INET, AF_INET6
@@ -337,21 +337,6 @@
 type[1] = 'R'
 type[2] = 'L'
 
-# from include/net/tcp_states.h:
-tcpstate = {}
-tcpstate[1] = 'ESTABLISHED'
-tcpstate[2] = 'SYN_SENT'
-tcpstate[3] = 'SYN_RECV'
-tcpstate[4] = 'FIN_WAIT1'
-tcpstate[5] = 'FIN_WAIT2'
-tcpstate[6] = 'TIME_WAIT'
-tcpstate[7] = 'CLOSE'
-tcpstate[8] = 'CLOSE_WAIT'
-tcpstate[9] = 'LAST_ACK'
-tcpstate[10] = 'LISTEN'
-tcpstate[11] = 'CLOSING'
-tcpstate[12] = 'NEW_SYN_RECV'
-
 # process event
 def print_ipv4_event(cpu, data, size):
     event = b["ipv4_events"].event(data)
@@ -362,9 +347,9 @@ def print_ipv4_event(cpu, data, size):
         "%s:%s" % (inet_ntop(AF_INET, pack('I', event.daddr)), event.dport)),
         end='')
     if args.sequence:
-        print(" %-12s %s" % (tcpstate[event.state], event.seq))
+        print(" %-12s %s" % (tcp.state2str(event.state), event.seq))
     else:
-        print(" %s" % (tcpstate[event.state]))
+        print(" %s" % (tcp.state2str(event.state)))
 
 def print_ipv6_event(cpu, data, size):
     event = b["ipv6_events"].event(data)
@@ -375,9 +360,9 @@ def print_ipv6_event(cpu, data, size):
         "%s:%d" % (inet_ntop(AF_INET6, event.daddr), event.dport)),
         end='')
     if args.sequence:
-        print(" %-12s %s" % (tcpstate[event.state], event.seq))
+        print(" %-12s %s" % (tcp.state2str(event.state), event.seq))
     else:
-        print(" %s" % (tcpstate[event.state]))
+        print(" %s" % (tcp.state2str(event.state)))
 
 def depict_cnt(counts_tab, l3prot='ipv4'):
     for k, v in sorted(counts_tab.items(), key=lambda counts: counts[1].value):

tools/tcpstates.py

@@ -16,7 +16,7 @@
 # 20-Mar-2018   Brendan Gregg   Created this.
 
 from __future__ import print_function
-from bcc import BPF
+from bcc import BPF, tcp
 import argparse
 from socket import inet_ntop, AF_INET, AF_INET6
 import sys
@@ -312,28 +312,6 @@
         exit(1)
 
 
-def tcpstate2str(state):
-    # from include/net/tcp_states.h:
-    tcpstate = {
-        1: "ESTABLISHED",
-        2: "SYN_SENT",
-        3: "SYN_RECV",
-        4: "FIN_WAIT1",
-        5: "FIN_WAIT2",
-        6: "TIME_WAIT",
-        7: "CLOSE",
-        8: "CLOSE_WAIT",
-        9: "LAST_ACK",
-        10: "LISTEN",
-        11: "CLOSING",
-        12: "NEW_SYN_RECV",
-    }
-
-    if state in tcpstate:
-        return tcpstate[state]
-    else:
-        return str(state)
-
 def journal_fields(event, addr_family):
     addr_pfx = 'IPV4'
     if addr_family == AF_INET6:
@@ -354,8 +332,8 @@ def journal_fields(event, addr_family):
         'OBJECT_TCP_SOURCE_PORT': str(event.lport),
         'OBJECT_' + addr_pfx + '_DESTINATION_ADDRESS': inet_ntop(addr_family, event.daddr),
         'OBJECT_TCP_DESTINATION_PORT': str(event.dport),
-        'OBJECT_TCP_OLD_STATE': tcpstate2str(event.oldstate),
-        'OBJECT_TCP_NEW_STATE': tcpstate2str(event.newstate),
+        'OBJECT_TCP_OLD_STATE': tcp.state2str(event.oldstate),
+        'OBJECT_TCP_NEW_STATE': tcp.state2str(event.newstate),
         'OBJECT_TCP_SPAN_TIME': str(event.span_us)
         }
 
@@ -396,7 +374,7 @@ def print_event(event, addr_family):
         version if args.wide or args.csv else "",
         inet_ntop(addr_family, event.saddr), event.lport,
         inet_ntop(addr_family, event.daddr), event.dport,
-        tcpstate2str(event.oldstate), tcpstate2str(event.newstate),
+        tcp.state2str(event.oldstate), tcp.state2str(event.newstate),
         float(event.span_us) / 1000))
     if args.journal:
         journal.send(**journal_fields(event, addr_family))