From 8ce56e19bedf1f781f69d81b5cd89843aff519b7 Mon Sep 17 00:00:00 2001
From: Thomas Williams <thomas@ycwmwl.cymru>
Date: Mon, 19 Aug 2024 22:44:26 +0100
Subject: [PATCH] Some further changes to assist with bottleknecking causing
 upstream issues to UFW

---
 reportIPBlock.py | 97 +++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 79 insertions(+), 18 deletions(-)

diff --git a/reportIPBlock.py b/reportIPBlock.py
index 37b46eb..0f2fe35 100644
--- a/reportIPBlock.py
+++ b/reportIPBlock.py
@@ -1,33 +1,94 @@
 def main():
 
     import socket
+    import os
     import config
     import datetime
+    import signal
+    import threading
     from log import logsManager
     import sys
-    
-    if len(sys.argv) != 4:
-        print("Usage: python script.py <blockedIPAddress> <jail> <live>")
-        sys.exit(1)
+    import select
 
-    hostname = socket.gethostname()
-    ipAddress = socket.gethostbyname(hostname)
-    blockedIPAddress = sys.argv[1]
-    jail = sys.argv[2]
-    live = int(sys.argv[3])
-    logTime = datetime.datetime.now()
-    
-    if not config.loggingMode == 'none' and not config.loggingMode == 'rabbitmq':
+    socketFile = '/etc/monutil/ip.socket'
 
-        manager = logsManager(config.sqlServer, config.sqlDatabase, config.sqlUsername, config.sqlPassword)
-        manager.insertIPBlock(hostname, ipAddress, blockedIPAddress, jail, live, logTime, token)
+    if len(sys.argv) == 4:
 
-    if config.loggingMode == 'rabbitmq':
+        hostname = socket.gethostname()
+        ipAddress = socket.gethostbyname(hostname)
+        blockedIPAddress = sys.argv[1]
+        jail = sys.argv[2]
+        live = int(sys.argv[3])
+        logTime = datetime.datetime.now()
+
+        clientSocket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        clientSocket.connect(socketFile)
+        data = f"{hostname}|{ipAddress}|{logTime}|ipBlock|{blockedIPAddress}|{jail}|{live}"
+        clientSocket.send(data.encode('utf-8'))
+        clientSocket.close()
+        sys.exit(0)
+
+    else:
         
-        import rabbitmq
+        if os.path.exists(socketFile):
+            os.remove(socketFile)
+
+        stop_event = threading.Event()
+        dataBuffer = []
+
+        def publishData(stop_event):
+
+            if config.loggingMode = 'rabbitmq':
+
+                import rabbitmq
+                rabbitmq = rabbitmq.rabbitMQClient(config.rabbitmqca,config.rabbitmqcacert,config.rabbitmqcakey,config.rabbitmqHost,config.rabbitmqPort,config.rabbitmqRoutingKey)
+            
+            while not (stop_event.is_set()):
+                if dataBuffer:
+                    data = dataBuffer.pop(0)
+                    
+                    if config.loggingMode = 'rabbitmq':
+                        rabbitmq.publish(f"{data}")
+                    else:
+                        print("Not yet implemented")
+
+        def cleanup(signum, frame):
+            
+            print("Signal received, shutting down...")
+            stop_event.set() 
+
+        signal.signal(signal.SIGTERM, cleanup)
+        signal.signal(signal.SIGINT, cleanup)
+
+        def server(stop_event):
+
+            serverSocket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            serverSocket.bind(socketFile)
+            serverSocket.listen(50)
+
+            while not (stop_event.is_set()):
+            
+                readable, _, _ = select.select([serverSocket], [], [], 1.0)
+                
+                if readable:
+                    clientSocket, _ = serverSocket.accept()
+                    data = clientSocket.recv(1024).decode('utf-8')
+                    
+                    if data:
+                
+                        dataBuffer.append(data)
+                        print(data)
+            
+                    clientSocket.close()
+
+        publishThread = threading.Thread(target=publishData, args=(stop_event,))
+        serverThread = threading.Thread(target=server, args=(stop_event,))
         
-        rabbitmq = rabbitmq.rabbitMQClient(config.rabbitmqca,config.rabbitmqcacert,config.rabbitmqcakey,config.rabbitmqHost,config.rabbitmqPort,config.rabbitmqRoutingKey)
-        rabbitmq.publish(hostname + '|' + socket.gethostbyname(socket.gethostname()) + '|' + str(logTime) + '|' + 'ipBlock' + '|' + str(blockedIPAddress) + '|' + str(jail) + '|' + str(live))
+        publishThread.start()
+        serverThread.start()
+        
+        publishThread.join()
+        serverThread.join()
 
 if __name__ == "__main__":
     main()