Update main.py

Author: aniko33

server/main.py

@@ -1,166 +1,229 @@
-import rsa
+from threading import Thread
 import socket
+import rsa
 import zlib
+import sys
 import time
 import json
-import _thread as thd
 
 from rich import print
 
-#start 
+# Read config file
 config_file = "config.json"
 with open(config_file, "r") as f:
     config_json = json.load(f)
 
-#vars
+# Vars
 clients = []
-username = {}
+nicknames = []
+
+# Config
 
-#config
-ip = config_json["ip"]
-port = config_json["port"]
-buffer = config_json["buffer"]
-welcome_message = config_json["welcome_message"]
+"""
+ip : The IP address where the server will start listening for connections
+port : The connection port, it is recommended to keep it default (8889)
+buffer : The maximum network buffer
+welcome_message : Welcome message to new users
+"""
 
-#create socket
+ip: str = config_json["ip"]
+port: int = config_json["port"]
+buffer: int = config_json["buffer"]
+welcome_message: str = config_json["welcome_message"]
+
+# Create socket
 server = socket.socket()
 
-#start listing
-server.bind((ip,port))
+# Start listing
+server.bind((ip, port))
 server.listen(32)
 
-#creating RSA key
-def create_keys(buffer: int):
-    public_key, private_key = rsa.newkeys(buffer)
-    return public_key, private_key
+
+class API:
+
+    """
+    def create_keys (buffer: int):
+        Generate an RSA key
+
+    def send_buffer (socket, buffer: int):
+        Send the buffer to the client
+
+    class Chat:
+        arg: private_key, public_key
+
+        def send (socket, message: str):
+            Send encrypted message
+
+        def recv (socketm, message: str):
+            Receives a message and decrypt it
+
+    class RSA:
+        arg: public_key, private_key
+
+        def encryption (message: str):
+            Encrypt message
+
+        def decrypt (message: bytes):
+            Decrypt message
+    """
+
+    def create_keys(buffer: int):
+        public_key, private_key = rsa.newkeys(buffer)
+        return public_key, private_key
+
+    def send_buffer(s, buffer: int):
+        s.send(str(buffer).encode())
+
+    class Chat:
+        def __init__(self, priv_key, pub_key) -> None:
+            self.priv_key = priv_key
+            self.pub_key = pub_key
+
+        def send(self, s, msg: str):
+            s.send(rsa.encrypt(msg.encode(), self.pub_key))
+
+        def recv(self, s, buffer: int):
+            msg = s.recv(buffer)
+            return rsa.decrypt(msg, self.priv_key)
+
+    class Send_keys:
+        def __init__(self, pub_key, priv_key, client) -> None:
+            self.client = client
+            self.pub_key = pub_key
+            self.priv_key = priv_key
+
+        def private(self):
+            private_key_exported = rsa.PrivateKey.save_pkcs1(self.priv_key)
+            # compressing
+            private_key_exported = zlib.compress(private_key_exported, 4)
+            self.client.send(private_key_exported)
+
+        def public(self):
+            public_key_exported = rsa.PublicKey.save_pkcs1(self.pub_key)
+            # compressing
+            public_key_exported = zlib.compress(public_key_exported, 4)
+            self.client.send(public_key_exported)
+
+    class RSA:
+        def __init__(self, pub_key, priv_key) -> None:
+            self.pub_key = pub_key
+            self.priv_key = priv_key
+
+        def encrypt(self, msg: str):
+            return rsa.encrypt(msg.encode(), self.pub_key)
+
+        def decrypt(self, msg: bytes):
+            return rsa.decrypt(msg, self.priv_key)
+
 
 class Chat:
-    def __init__(self, priv_key, pub_key) -> None:
-        self.priv_key = priv_key
-        self.pub_key = pub_key
-    def send(self, s, msg:str):
-        s.send(rsa.encrypt(msg.encode(), self.pub_key))
-    def recv(self, s, buffer: int):
-        msg = s.recv(buffer)
-        return rsa.decrypt(msg, self.priv_key)
-
-class Send_keys:
-    def __init__(self,pub_key, priv_key, client) -> None:
+    """
+    Args: client (socket), private_key, public_key
+
+    def joined (nickname: str):
+        It will send a message when a client disconnect
+
+    def welcome_message (bytes: bytes):
+        It will send the clients the encrypted welcome message
+
+    def send_to_clients (message: bytes):
+        It sends clients a message, but it won't be able to send it to itself
+
+    def remove_client (client):
+        removes clients from the client list
+
+    def middle:
+        When a customer enters the chat, perform this function.
+        Send clients a message announcing that a client has logged in,
+        then wait for a message from the client and then send it to all
+
+    def run:
+        It's where this class will launch
+    """
+
+    def __init__(self, client, private_key, public_key) -> None:
         self.client = client
-        self.pub_key = pub_key
-        self.priv_key = priv_key
+        self.private_key = private_key
+        self.public_key = public_key
 
-    def private(self):
-        private_key_exported = rsa.PrivateKey.save_pkcs1(self.priv_key)
-        #compressing
-        private_key_exported = zlib.compress(private_key_exported, 4)
-        self.client.send(private_key_exported)
+    def joined(self, nickname: str):
+        self.send_to_clients(self.rsa_api.encrypt(
+            f"[green]{nickname}[/green] has joined."))
 
-    def public(self):
-        public_key_exported = rsa.PublicKey.save_pkcs1(self.pub_key)
-        #compressing
-        public_key_exported = zlib.compress(public_key_exported, 4)
-        self.client.send(public_key_exported)
+    def welcome_message(self, welcome_message: bytes):
+        self.client.send(welcome_message)
 
-class Main:
-    def __init__(self) -> None:
-        pass
-
-    def chat(client, public_key, private_key):
-        def send_buffer():
-            client.send(str(buffer).encode())
-
-        def remove_connection(client_socket):
-            if client_socket in clients:
-                #remove client socket
-                clients.remove(client_socket)
-                #remove username
-                username.pop(client_socket)
-                print("[[yellow]?[/yellow]] Client disconnected")
-
-        def send_to_clients(msg):
-            for client_socket in clients:
-                if client != client_socket:
-                    try:
-                        client_socket.send(msg)
-                    except:
-                        client_socket.close()
-                        remove_connection(client_socket)
-
-        def check_username():
-            exist = False
-            
-            username_unchecked = client.recv(buffer).decode()
-            for user in username:
-                if not exist:
-                    if username[user] == username_unchecked:
-                        exist = True
-            #send code status and close connection
-            if exist:
-                client.send("False".encode())
-                remove_connection(client)
-            #send code status and add username
-            else: 
-                client.send("True".encode())
-                username[client] = username_unchecked
-        def middle():
-            while True:
+    def send_to_clients(self, msg: bytes):
+        for client in clients:
+            if client != self.client:
                 try:
-                    msg = client.recv(buffer)
-                    if msg.decode() == "/users":
-                        users = ""
-                        for user in username:
-                            users += username[user]+","
-                        client.send(users.encode())
-                    else:
-                        if msg:
-                            send_to_clients(msg)
-                        else:
-                            remove_connection(client)
-                except:
-                    continue
-        
-        send_buffer()
-        
-        #init class for send RSA keys
-        client_key = Send_keys(public_key, private_key, client)
-        
-        #send keys
-        client_key.private(); time.sleep(0.5)
-        client_key.public()
-        
-        #API
-        chat_sock = Chat(private_key, public_key)
+                    client.send(msg)
+                except BaseException:
+                    self.remove_client(client)
+
+    def remove_client(self, client):
+        print("[[yellow]?[/yellow]] Client disconnected")
+        index = clients.index(client)
+        clients.remove(client)
+        nickname = nicknames[index]
+        self.send_to_clients(self.rsa_api.encrypt(
+            f"[green]{nickname}[/green] has left."))
+        nicknames.remove(nickname)
+
+    def middle(self):
+        index = clients.index(self.client)
+        nickname = nicknames[index]
+        self.joined(nickname)
+        while True:
+            try:
+                msg = self.client.recv(buffer)
+                self.send_to_clients(msg)
+            except BaseException:
+                self.remove_client(self.client)
+                break
 
-        time.sleep(0.5)
+    def run(self):
+        API.send_buffer(self.client, buffer)
+        send_keys = API.Send_keys(
+            self.public_key,
+            self.private_key,
+            self.client)
 
-        #check username
-        check_username()
-        
+        self.rsa_api = API.RSA(self.public_key, self.private_key)
+        self.chat_api = API.Chat(self.private_key, self.public_key)
+
+        send_keys.public()
         time.sleep(0.5)
+        send_keys.private()
+        # Encrypt welcome_message and send to client
+        self.welcome_message(self.rsa_api.encrypt(welcome_message))
+        self.middle()
 
-        #welcome message
-        chat_sock.send(client, welcome_message)
 
-        middle()
-    def run(self):
-        stopped = False
-        
-        print(f"[[cyan]+[/cyan]] Buffer: {buffer}")
-
-        print("[[cyan]+[/cyan]] Creating RSA keys...")
-        self.public_key, self.private_key = create_keys(buffer)
-        print("[[cyan]+[/cyan]] RSA keys created")
-        
-        while (not stopped):
+class Main:
+    """
+    def run:
+        It will generate keys and wait for connections
+    """
+    def run():
+        print(f"[[magenta]*[/magenta]] Buffer: {buffer}")
+        print("[[cyan]+[/cyan]] RSA key generation...")
+        public_key, private_key = API.create_keys(buffer)
+        print("[[cyan]+[/cyan]] RSA key generated")
+        while True:
             client, addr = server.accept()
             print("[[yellow]?[/yellow]] Client connected")
-            #add 
+
+            nickname = client.recv(buffer).decode()
+            nicknames.append(nickname)
+
             clients.append(client)
 
-            #start in idle
-            thd.start_new_thread(Main.chat, (client, self.public_key, self.private_key))
+            chat = Chat(client, private_key, public_key)
+
+            multi_conn = Thread(target=chat.run)
+            multi_conn.start()
+
 
 if __name__ == "__main__":
-    main_start = Main()
-    main_start.run()
+    Main.run()