diff --git a/server/main.py b/server/main.py index e884dce..dc660db 100644 --- a/server/main.py +++ b/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()