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 server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # set the socket options, SO_REUSEADDR allows to reuse the same port # SO_KEEPALIVE enables the socket to send keepalive messages server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) # bind the socket to the address and port server_socket.bind(('', CHATPORT)) # listen for incoming connections server_socket.listen(MAX_CONNECT_QUEUE) # set the server socket to non-blocking server_socket.setblocking(False) # create a list to store the sockets, server and client socket_list = [server_socket] # 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}") # infinite loop for the program while quitFlag: # use the select function to monitor the sockets read_sockets, _, _ = select.select(socket_list, [], []) # loop through the sockets for sock in read_sockets: # if the socket is the server socket, accept the connection if sock == server_socket: client_socket, client_address = server_socket.accept() client_socket.setblocking(False) socket_list.append(client_socket) print("\nAccepted connection from new client.") client_socket.send(b"\nWelcome to the chat server!\n") 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 try: message = sock.recv(BUFFER_SIZE).decode('utf-8') if message: state = client_states.get(sock, "NORMAL") if state == "WAITING_FOR_USERNAME": 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": client_info[sock]["username"] = message.strip() sock.send(b"\nUsername changed successfully.\n") 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] chat_rooms[sock] = message.strip() sock.send(f"\nSwitched to chatroom '{chat_rooms[sock]}.'\n".encode('utf-8')) print(f"\nClient switched to chatroom '{chat_rooms[sock]}.'") if old_chatroom != "default" and old_chatroom not in chat_rooms.values(): print(f"\nChatroom '{old_chatroom}' is empty and removed.") client_states[sock] = "NORMAL" else: if message.strip() == "%!": sock.close() socket_list.remove(sock) del client_info[sock] del chat_rooms[sock] del client_states[sock] print("\nClient has closed connection.") elif message.strip() == "%?": print("\nClient requested help.") sock.send(b"Help: (%?)\n") sock.send(b"Send messages like normal. Use enter to send a message.\n") sock.send(b"Available Commands:\n") sock.send(b"\t%? : Show this help prompt.\n") 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 = client_info[sock]["username"] chat_room = chat_rooms[sock] send_message = f"{username} >> {message}".encode('utf-8') for client_sock in socket_list: if client_sock != server_socket and chat_rooms.get(client_sock) == chat_room: client_sock.send(send_message) else: sock.close() socket_list.remove(sock) if sock in client_info: del client_info[sock] if sock in chat_rooms: old_chatroom = chat_rooms[sock] del chat_rooms[sock] if old_chatroom != "default" and old_chatroom not in chat_rooms.values(): print(f"\nChatroom '{old_chatroom}' is empty and has been removed.") if sock in client_states: del client_states[sock] print("\nClient has closed connection.") except Exception as err: print(f"Error: {err}") sock.close() socket_list.remove(sock) if sock in client_info: del client_info[sock] if sock in chat_rooms: old_chatroom = chat_rooms[sock] del chat_rooms[sock] if old_chatroom != "default" and old_chatroom not in chat_rooms.values(): print(f"\nChatroom '{old_chatroom}' is empty and has been removed.") if sock in client_states: 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__": main()