implement a NbiContainer class to organize NBIs effectively

fakenet/diverters/diverterbase.py

@@ -12,6 +12,7 @@
 import subprocess
 from . import fnpacket
 from . import fnconfig
+from . import nbicontainer
 from .debuglevels import *
 from collections import namedtuple
 from collections import OrderedDict
@@ -1796,21 +1797,26 @@ def maybeExecuteCmd(self, pkt, pid, comm):
     def mapOrigSportToProxySport(self, orig_sport, proxy_sport):
         self.proxy_original_source_ports[proxy_sport] = orig_sport
 
-    def logNbi(self, listener_port, nbi):
+    def logNbi(self, listener_name, listener_port, nbi):
         proxied_nbi = listener_port in self.proxy_original_source_ports
         orig_source_port = self.proxy_original_source_ports[listener_port] if proxied_nbi else listener_port
         _, __, pid, comm = self.sessions[orig_source_port]
 
         # If it's a new NBI from an exisitng process, append nbi, else create new key
         existing_process = (pid, comm) in self.nbi
-        if existing_process:
-            # {(123, chrome.exe): {"host": ["www.google.com"], "version": ["HTTP1.1"]}}
-            for nbi_attributes in nbi:
-                if nbi_attributes in self.nbi[(pid, comm)]:
-                    self.nbi[(pid, comm)][nbi_attributes].append(nbi[nbi_attributes][0])
-                else:
-                    self.nbi[(pid, comm)][nbi_attributes] = nbi[nbi_attributes]
+
+        import pdb;pdb.set_trace()
+
+        if not existing_process:
+            nbiValue = nbicontainer.NbiContainer()
+            nbiValue.listenerType(listener_name)
+            nbiValue.addNbi(orig_source_port, nbi)
+            self.nbi[(pid, comm)] = nbiValue
 
         else:
-            self.nbi[(pid, comm)] = nbi
+            self.nbi[(pid, comm)].listenerType(listener_name)
+            self.nbi[(pid, comm)].addNbi(orig_source_port, nbi)
 
+        for processes in self.nbi:
+            self.nbi[processes].prettyPrintNbi(processes)
+

fakenet/diverters/nbicontainer.py

@@ -0,0 +1,44 @@
+import logging
+
+INDENT = '  '
+
+class NbiContainer:
+
+    def __init__(self):
+        self.nbis = {}
+        self.listener = None
+        self.logger = logging.getLogger('NbiContainer')
+        self.logger.setLevel(logging.DEBUG)
+
+    def listenerType(self, listener):
+        if listener.lower()=="rawlistener":
+            self.listener = "TCP"
+        elif listener.lower()=="httplistener":
+            self.listener = "HTTP"
+
+    def addNbi(self, sport, nbi):
+        if self.listener == "HTTP":
+            # our nbi is a http request
+            self.nbis[sport] = nbi
+        elif self.listener == "TCP":
+            pass
+
+    def prettyPrintNbi(self, process_info):
+        if not self.nbis:
+            self.logger.info("No NBIs to log")
+        else:
+            pid = process_info[0]
+            comm = process_info[1]
+            if self.listener == "HTTP":
+                self.logger.info("Logging HTTP NBIs")
+                self.logger.info("Process ID: %s\t Process Name: %s", pid, comm)
+                for sport in self.nbis:
+                    self.logger.info(sport)
+                    for attributes in self.nbis[sport]:
+                        self.logger.info(INDENT + "%s: %s", attributes, self.nbis[sport][attributes])
+            elif self.listener == "TCP":
+                pass
+
+
+
+

fakenet/listeners/HTTPListener.py

@@ -366,19 +366,19 @@ def do_POST(self):
     def collect_nbi(self, requestline, headers, post_data=None):
         nbi = {}
         method, uri, version = requestline.split(" ")
-        nbi["method"] = [method]
-        nbi["uri"] = [uri]
-        nbi["version"] = [version]
+        nbi["method"] = method
+        nbi["uri"] = uri
+        nbi["version"] = version
 
         for line in str(headers).rstrip().split("\n"):
             key, _, value = line.partition(":")
-            nbi[key] = [value.lstrip()]
+            nbi[key] = value.lstrip()
 
         if post_data:
-            nbi["post_data"] = [post_data]
+            nbi["post_data"] = post_data
 
         # report diverter each time when NBI is reported
-        self.server.diverter.logNbi(self.client_address[1], nbi)
+        self.server.diverter.logNbi('HTTPListener', self.client_address[1], nbi)
 
     def get_response(self, path):
         response = "<html><head><title>FakeNet</title><body><h1>FakeNet</h1></body></html>"