Add lots of functionality to the server codebase.

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, 
+    # load persistent data
-    # and each one can only be associated one way at once, use dictionaries to
+    client_info, chat_rooms, client_states = load_data()
-    # associate with each other.
+    uuid_dict = load_uuids()
-    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 = {}
     # 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()
+                            username = message.strip()
-                            chat_rooms[sock] = "default"
+                            if username == ADMIN_USERNAME:
-                            sock.send(b"\nUsername accepted. You are now in the 'default' chatroom.\n")
+                                sock.send(b"\nEnter admin password: \n")
-                            client_states[sock] = "NORMAL"
+                                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%@ <username> <message> : 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: %@ <username> <message>\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:
+                        if sock in client_info:
-                            del usernames[sock]
+                            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:
+                    if sock in client_info:
-                        del usernames[sock]
+                        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__":