#!/usr/bin/python3

# MIT License

# Copyright (c) 2024 Thomas Williams - https://git.server.wales/thomas

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

def main():

    import socket
    import os
    import config
    import datetime
    import signal
    import threading
    import log
    import sys
    import select
    import time

    socketFile = '/etc/monutil/ip.socket'

    if len(sys.argv) == 4:

        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:
        
        if os.path.exists(socketFile):
            os.remove(socketFile)

        stopEvent = threading.Event()
        failureEvent = threading.Event()
        
        dataBuffer = []

        def publishData(stopEvent):

            try:

                if config.loggingMode == 'rabbitmq':

                    import rabbitmq
                    rabbitmq = rabbitmq.rabbitMQClient(config.rabbitmqca,config.rabbitmqcacert,config.rabbitmqcakey,config.rabbitmqHost,config.rabbitmqPort,config.rabbitmqRoutingKey)
            
                    while not (stopEvent.is_set() and not failureEvent.is_set()):

                        time.sleep(1)
                        
                        if dataBuffer:
                            
                            data = dataBuffer.pop(0)
                            rabbitmq.publish(f"{data}")
                        
                else:

                    while not (stopEvent.is_set() and not failureEvent.is_set()):

                        time.sleep(1)

                        if dataBuffer:

                            data = dataBuffer.pop(0)
                            data = data.split('|')
                            manager = log.logsManager(config.sqlServer, config.sqlDatabase, config.sqlUsername, config.sqlPassword)
                            manager.insertIPBlock(data[0], data[1], data[4], data[5], data[6], data[2], config.ipinfoAPIToken)

            except Exception:
                failureEvent.set()

        def cleanup(signum, frame):
            
            print("Signal received, shutting down...")
            stopEvent.set() 

        signal.signal(signal.SIGTERM, cleanup)
        signal.signal(signal.SIGINT, cleanup)

        def server(stopEvent):

            try:
                
                serverSocket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                serverSocket.bind(socketFile)
                serverSocket.listen(50)

                while not (stopEvent.is_set() and not failureEvent.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()

            except Exception:
                
                failureEvent.set()

        publishThread = threading.Thread(target=publishData, args=(stopEvent,))
        serverThread = threading.Thread(target=server, args=(stopEvent,))
        
        publishThread.start()
        serverThread.start()
        
        publishThread.join()
        serverThread.join()

        if failureEvent.is_set():
            print("One of the threads failed. Terminating")
            sys.exit(1)

if __name__ == "__main__":
    main()