From a2adaa05a69f946c3ae96f3cc0fd25d561002d50 Mon Sep 17 00:00:00 2001 From: cdez Date: Wed, 12 Mar 2025 21:09:40 -0700 Subject: [PATCH] Add lots of functionality to the server codebase. --- chatserver/server.py | 128 ++++++++++++++++++++++++++++++++++++------- 1 file changed, 108 insertions(+), 20 deletions(-) diff --git a/chatserver/server.py b/chatserver/server.py index fd79b83..dd6bb99 100644 --- a/chatserver/server.py +++ b/chatserver/server.py @@ -1,11 +1,59 @@ import socket import select +import uuid +import pickle +import os CHATPORT = 8181 BUFFER_SIZE = 1024 MAX_CONNECT_QUEUE = 10 +DATA_FILE = 'server_data.pkl' +ADMIN_USERNAME = "admin" +ADMIN_PASSWORD = "adminpass" +UUID_FILE = 'uuids.pkl' + +def save_data(client_info, chat_rooms, client_states): + serializable_client_info = {info["user_id"]: info for sock, info in client_info.items()} + serializable_chat_rooms = {info["user_id"]: chat_rooms[sock] for sock, info in client_info.items()} + serializable_client_states = {info["user_id"]: client_states[sock] for sock, info in client_info.items()} + with open(DATA_FILE, 'wb') as f: + pickle.dump((serializable_client_info, serializable_chat_rooms, serializable_client_states), f) + +def load_data(): + if os.path.exists(DATA_FILE): + with open(DATA_FILE, 'rb') as f: + serializable_client_info, serializable_chat_rooms, serializable_client_states = pickle.load(f) + client_info = {uuid.UUID(user_id): info for user_id, info in serializable_client_info.items()} + chat_rooms = {uuid.UUID(user_id): room for user_id, room in serializable_chat_rooms.items()} + client_states = {uuid.UUID(user_id): state for user_id, state in serializable_client_states.items()} + return client_info, chat_rooms, client_states + return {}, {}, {} + +def save_uuids(uuid_dict): + with open(UUID_FILE, 'wb') as f: + pickle.dump(uuid_dict, f) + +def load_uuids(): + if os.path.exists(UUID_FILE): + with open(UUID_FILE, 'rb') as f: + return pickle.load(f) + return {} + +def get_client_uuid(client_socket, uuid_dict): + client_address = client_socket.getpeername() + client_key = f"{client_address[0]}_{client_address[1]}" + + if client_key in uuid_dict: + user_id = uuid_dict[client_key] + else: + user_id = str(uuid.uuid4()) + uuid_dict[client_key] = user_id + save_uuids(uuid_dict) + + return user_id def main(): + global quitFlag # set the flag variable to true for infinite loop quitFlag = True # create a socket object @@ -23,14 +71,9 @@ def main(): # create a list to store the sockets, server and client socket_list = [server_socket] - # since a socket is associated with a username and a chatroom, - # and each one can only be associated one way at once, use dictionaries to - # associate with each other. - usernames = {} - chat_rooms = {} - # keep track of the states of each client connected, and if they are - # technically in an 'unsafe' state like in change username or chatroom mode. - client_states = {} + # load persistent data + client_info, chat_rooms, client_states = load_data() + uuid_dict = load_uuids() # let the user know the server is listening on port CHATPORT. print(f"Chatserver listening on port: {CHATPORT}") @@ -50,6 +93,9 @@ def main(): client_socket.send(b"\nUse '%!' to exit, or '%?' for help.\n") client_socket.send(b"\nType in a username: \n") client_states[client_socket] = "WAITING_FOR_USERNAME" + # Get or generate a unique user ID for the client + user_id = get_client_uuid(client_socket, uuid_dict) + client_info[client_socket] = {"user_id": user_id} # otherwise the socket is a client socket. else: # try to receive the message from the client @@ -58,14 +104,30 @@ def main(): if message: state = client_states.get(sock, "NORMAL") if state == "WAITING_FOR_USERNAME": - usernames[sock] = message.strip() - chat_rooms[sock] = "default" - sock.send(b"\nUsername accepted. You are now in the 'default' chatroom.\n") - client_states[sock] = "NORMAL" + username = message.strip() + if username == ADMIN_USERNAME: + sock.send(b"\nEnter admin password: \n") + client_states[sock] = "WAITING_FOR_ADMIN_PASSWORD" + client_info[sock]["username"] = username + else: + client_info[sock]["username"] = username + chat_rooms[sock] = "default" + sock.send(b"\nUsername accepted. You are now in the 'default' chatroom.\n") + client_states[sock] = "NORMAL" + elif state == "WAITING_FOR_ADMIN_PASSWORD": + password = message.strip() + if password == ADMIN_PASSWORD: + sock.send(b"\nAdmin authenticated successfully.\n") + client_info[sock]["is_admin"] = True + chat_rooms[sock] = "default" + client_states[sock] = "NORMAL" + else: + sock.send(b"\nInvalid admin password. Please enter a username: \n") + client_states[sock] = "WAITING_FOR_USERNAME" elif state == "WAITING_FOR_NEW_USERNAME": - usernames[sock] = message.strip() + client_info[sock]["username"] = message.strip() sock.send(b"\nUsername changed successfully.\n") - print(f"\nClient switched username to {usernames[sock]}.") + print(f"\nClient switched username to {client_info[sock]['username']}.") client_states[sock] = "NORMAL" elif state == "WAITING_FOR_NEW_CHATROOM": old_chatroom = chat_rooms[sock] @@ -79,7 +141,7 @@ def main(): if message.strip() == "%!": sock.close() socket_list.remove(sock) - del usernames[sock] + del client_info[sock] del chat_rooms[sock] del client_states[sock] print("\nClient has closed connection.") @@ -92,14 +154,38 @@ def main(): sock.send(b"\t%! : Exit from the Chat Server.\n") sock.send(b"\t%* : Switch to another chatroom.\n") sock.send(b"\t%$ : Change Username.\n") + sock.send(b"\t%@ : Send a direct message to a user.\n") + if client_info[sock].get("is_admin"): + sock.send(b"\t%stop : Stop the server.\n") elif message.strip() == "%*": sock.send(b"\nEnter chatroom name: \n") client_states[sock] = "WAITING_FOR_NEW_CHATROOM" elif message.strip() == "%$": sock.send(b"\nEnter new username: \n") client_states[sock] = "WAITING_FOR_NEW_USERNAME" + elif message.strip().startswith("%@ "): + parts = message.strip().split(' ', 2) + if len(parts) < 3: + sock.send(b"\nInvalid direct message format. Use: %@ \n") + else: + target_username = parts[1] + dm_message = parts[2] + target_sock = None + for s, info in client_info.items(): + if info.get("username") == target_username: + target_sock = s + break + if target_sock: + sender_username = client_info[sock]["username"] + target_sock.send(f"DM from {sender_username}: {dm_message}\n".encode('utf-8')) + sock.send(b"\nDirect message sent.\n") + else: + sock.send(b"\nUser not found.\n") + elif message.strip() == "%stop" and client_info[sock].get("is_admin"): + print("\nAdmin has stopped the server.") + quitFlag = False else: - username = usernames[sock] + username = client_info[sock]["username"] chat_room = chat_rooms[sock] send_message = f"{username} >> {message}".encode('utf-8') for client_sock in socket_list: @@ -108,8 +194,8 @@ def main(): else: sock.close() socket_list.remove(sock) - if sock in usernames: - del usernames[sock] + if sock in client_info: + del client_info[sock] if sock in chat_rooms: old_chatroom = chat_rooms[sock] del chat_rooms[sock] @@ -122,8 +208,8 @@ def main(): print(f"Error: {err}") sock.close() socket_list.remove(sock) - if sock in usernames: - del usernames[sock] + if sock in client_info: + del client_info[sock] if sock in chat_rooms: old_chatroom = chat_rooms[sock] del chat_rooms[sock] @@ -133,6 +219,8 @@ def main(): del client_states[sock] print("\nClient has closed connection.") + # Save data before closing the server + save_data(client_info, chat_rooms, client_states) server_socket.close() if __name__ == "__main__":