From def22643f6d6c24824b1badb58abe867045cd91e Mon Sep 17 00:00:00 2001
From: Cole Deck <cole@deck.sh>
Date: Mon, 24 Oct 2022 22:35:35 -0500
Subject: [PATCH] fix UI data for blocked connections

---
 block.py    |  4 +++-
 ippigeon.py | 25 ++++++++++++++++++++-----
 panel.py    |  9 +++++----
 util.py     |  7 ++++---
 4 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/block.py b/block.py
index 4083cb8..ee80967 100644
--- a/block.py
+++ b/block.py
@@ -34,6 +34,7 @@ def block_conn(config, datafile, res):
     fprint("Searching block data")
     mydata = list()
     badapps = list()
+    badips = list()
     with open(find_data_file(datafile), newline='') as csvfile:
         csvreader = csv.reader(csvfile, delimiter=',', quotechar='|')
         
@@ -66,7 +67,8 @@ def block_conn(config, datafile, res):
             if srcip == badsrcip or destip == baddestip and not pid in badapps:
                 fprint("FLAG " + srcip + " " + destip + " " + str(pid))
                 badapps.append(pid)
+                badips.append(baddestip)
                 #fprint(badapps)
                 #fprint("FLAG " + srcip + " " + destip + " " + str(pid))
                 #kill(pid)
-    return badapps
+    return badapps, badips
diff --git a/ippigeon.py b/ippigeon.py
index 1f2741c..5dbd6ae 100644
--- a/ippigeon.py
+++ b/ippigeon.py
@@ -82,16 +82,28 @@ def login_done(res):
 def blockdata_done(res):
     global settings
     fprint("FINISHED downloading block data")
+    tmpkill = settings["kill"]
+    settings["kill"] = False
     #block_res = pool.apply_async(block.block_conn, (config, datafile, res, settings))
-    block_res = block.block_conn(config, datafile, res)
+    block_pids, block_ips = block.block_conn(config, datafile, res)
     tmplist = settings["badapps"]
-    for x in block_res:
+    
+    for x in block_pids:
         fprint(x)
         if not x in tmplist:
             tmplist.append(x)
-        
     settings["badapps"] = tmplist
     fprint(settings["badapps"])
+    
+    tmplist = settings["badips"]
+    for x in block_ips:
+        fprint(x)
+        if not x in tmplist:
+            tmplist.append(x)
+    settings["badips"] = tmplist
+    fprint(settings["badips"])
+
+    settings["kill"] = tmpkill
         
 def killall():
     kids = active_children()
@@ -152,10 +164,12 @@ def mainloop(pool):
         settings["block"] = False
 
     if settings["kill"] == True:
-        for x in settings["badapps"]:
-            kill(x)
+        tmplist = settings["badapps"]
         settings["badapps"] = list()
 
+        for x in tmplist:
+            kill(x)
+
     sleep(interval / (interval * config["core"]["clockspeed"]))
     counter += 1
     if counter == interval * config["core"]["clockspeed"]:
@@ -202,6 +216,7 @@ if __name__ == '__main__':
             settings["continueui"] = False
             settings["killbox"] = list()
             settings["badapps"] = list()
+            settings["badips"] = list()
             settings["block"] = False
             settings["kill"] = False
             killme = manager.Value('d', 0)
diff --git a/panel.py b/panel.py
index c2ca5aa..cb8b19e 100644
--- a/panel.py
+++ b/panel.py
@@ -187,11 +187,12 @@ class ServerPanel(wx.Panel):
         else:
             fprint("Unable to delete")
         for i in range(len(TEST_FILE)):
+            
             idx = 0
-            for pid in settings["badapps"]:
+            for ip in settings["badips"]:
                 #fprint(pid)
                 idx = i
-                if str(TEST_FILE.iloc[i, 5]) == str(pid): # "bad" pid, highlight in red at the top
+                if str(TEST_FILE.iloc[i, 3]).find(ip) >= 0 and str(TEST_FILE.iloc[i, 4]).find("TIME_WAIT") < 0: # "bad" pid, highlight in red at the top
                     idx = self.list_ctrl.InsertItem(0, TEST_FILE.iloc[i, 0])
                     self.list_ctrl.SetItemBackgroundColour(idx, wx.Colour(200, 51, 51))
                     break
@@ -202,8 +203,8 @@ class ServerPanel(wx.Panel):
             
                     #fprint("Got " + TEST_FILE.iloc[i, 5])
             for j in range(1, 6):
-                fprint(str(idx) + " " + str(TEST_FILE.iloc[i, 0]))
-                self.list_ctrl.SetItem(idx, j, str(TEST_FILE.iloc[idx, j]))
+                #fprint(str(idx) + " " + str(TEST_FILE.iloc[i, 0]))
+                self.list_ctrl.SetItem(idx, j, str(TEST_FILE.iloc[i, j]))
 
                 #fprint(i, j, TEST_FILE.iloc[i, j])
         #self.SetSizer(self.main_sizer)
diff --git a/util.py b/util.py
index fbf7de3..0c54e03 100644
--- a/util.py
+++ b/util.py
@@ -42,9 +42,10 @@ def time():
 
 def kill(pid):
     setup_child()
-    fprint("Killing PID " + str(pid))
-    os.kill(int(pid), 9)
-    fprint("Signal 9 sent to PID " + str(pid))
+    if pid > 0:
+        fprint("Killing PID " + str(pid))
+        os.kill(int(pid), 9)
+        fprint("Signal 9 sent to PID " + str(pid))
 
 def fprint(msg):
     #if not getattr(sys, "frozen", False):