Share TCP-related functions across plugins (Fixes: #54)

mami-project · Aug 8, 2016 · 9903ea7 · 9903ea7
1 parent a06f3d1
commit 9903ea7
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 14 deletions.
diff --git a/pathspider/observer.py → pathspider/observer/__init__.py b/pathspider/observer.py → pathspider/observer/__init__.py
diff --git a/pathspider/observer/tcp.py b/pathspider/observer/tcp.py
@@ -0,0 +1,22 @@
+
+def tcp_setup(rec, ip):
+    rec['fwd_fin'] = False
+    rec['fwd_rst'] = False
+    rec['rev_fin'] = False
+    rec['rev_rst'] = False
+
+    return True
+
+def tcp_complete(rec, tcp, rev): # pylint: disable=W0612,W0613
+    if tcp.fin_flag and rev:
+        rec['rev_fin'] = True
+    if tcp.fin_flag and not rev:
+        rec['fwd_fin'] = True
+    if tcp.rst_flag and rev:
+        rec['rev_rst'] = True
+    if tcp.rst_flag and not rev:
+        rec['fwd_rst'] = True
+
+    return not ( ( rec['fwd_fin'] and rec['rev_fin'] ) or
+                 rec['fwd_rst'] or rec['rev_rst'] )
+
diff --git a/pathspider/plugins/dscpspider.py b/pathspider/plugins/dscpspider.py
@@ -14,6 +14,9 @@
 from pathspider.observer import basic_flow
 from pathspider.observer import basic_count
 
+from pathspider.observer.tcp import tcp_setup
+from pathspider.observer.tcp import tcp_complete
+
 Connection = collections.namedtuple("Connection", ["client", "port", "state"])
 SpiderRecord = collections.namedtuple("SpiderRecord", ["ip", "rport", "port",
                                                        "host", "dscp",
@@ -25,15 +28,6 @@
 
 ## Chain functions
 
-def tcp_setup(rec, ip):
-    rec['fin_count'] = 0
-    return True
-
-def tcp_complete(rec, tcp, rev): # pylint: disable=W0612,W0613
-    if tcp.fin_flag:
-        rec['fin_count'] += 1
-    return rec['fin_count'] < 2
-
 def dscp_setup(rec, ip):
     rec['fwd_dscp'] = None
     rec['rev_dscp'] = None

diff --git a/pathspider/plugins/ecnspider3.py b/pathspider/plugins/ecnspider3.py
@@ -14,6 +14,9 @@
 from pathspider.observer import basic_flow
 from pathspider.observer import basic_count
 
+from pathspider.observer.tcp import tcp_setup
+from pathspider.observer.tcp import tcp_complete
+
 Connection = collections.namedtuple("Connection", ["client", "port", "state"])
 SpiderRecord = collections.namedtuple("SpiderRecord", ["ip", "rport", "port",
                                                        "host", "ecnstate",
@@ -27,9 +30,6 @@
 
 ## Chain functions
 
-def tcpcompleted(rec, tcp, rev): # pylint: disable=W0612,W0613
-    return not tcp.fin_flag
-
 def ecnsetup(rec, ip):
     rec['ecn_zero'] = False
     rec['ecn_one'] = False
@@ -146,10 +146,10 @@ def create_observer(self):
         logger.info("Creating observer")
         try:
             return Observer(self.libtrace_uri,
-                            new_flow_chain=[basic_flow, ecnsetup],
+                            new_flow_chain=[basic_flow, tcp_setup, ecnsetup],
                             ip4_chain=[basic_count, ecncode],
                             ip6_chain=[basic_count, ecncode],
-                            tcp_chain=[ecnflags, tcpcompleted])
+                            tcp_chain=[ecnflags, tcp_complete])
         except:
             logger.error("Observer not cooperating, abandon ship")
             traceback.print_exc()