From 38c96879afd5f6e51c58e0b1cf8537b864e2fad4 Mon Sep 17 00:00:00 2001 From: tcsenpai Date: Wed, 19 Feb 2025 11:18:33 +0100 Subject: [PATCH] multiroom support --- data/rooms.json | 5 +++ libs/broadcast.py | 27 +++++++------ libs/process_message.py | 53 ++++++++++++++++++++++--- libs/room_manager.py | 86 +++++++++++++++++++++++++++++++++++++++++ main.py | 30 ++++++++++++-- 5 files changed, 180 insertions(+), 21 deletions(-) create mode 100644 data/rooms.json create mode 100644 libs/room_manager.py diff --git a/data/rooms.json b/data/rooms.json new file mode 100644 index 0000000..071017a --- /dev/null +++ b/data/rooms.json @@ -0,0 +1,5 @@ +{ + "lounge": { + "description": "The default chat room" + } +} \ No newline at end of file diff --git a/libs/broadcast.py b/libs/broadcast.py index 63b089c..4a62ed2 100644 --- a/libs/broadcast.py +++ b/libs/broadcast.py @@ -1,21 +1,24 @@ -def broadcast_message(connections, message, sender_addr=None): - """ - Broadcasts a message to all connected clients. - - Args: - connections (dict): Dictionary of active connections {addr: socket} - message (str): Message to broadcast - sender_addr (tuple, optional): Address of the sender to exclude - """ +def broadcast_message( + connections, message, sender_addr=None, room=None, system_msg=False +): + """Broadcasts a message to all users in a room or system-wide.""" formatted_message = f"\r\n{message}\r\n" encoded_message = formatted_message.encode("ascii") - for addr, conn in connections.items(): + # Get recipients based on room, unless it's a system message + recipients = ( + connections.keys() + if system_msg + else (room.users if room else connections.keys()) + ) + + for addr in recipients: + if addr not in connections: + continue try: - # Don't send back to the sender if specified if sender_addr and addr == sender_addr: continue - conn.sendall(encoded_message) + connections[addr].sendall(encoded_message) except (ConnectionError, BrokenPipeError): print(f"Error sending message to {addr}") except: diff --git a/libs/process_message.py b/libs/process_message.py index f22ac57..65f9033 100644 --- a/libs/process_message.py +++ b/libs/process_message.py @@ -2,8 +2,9 @@ from typing import Tuple class CommandProcessor: - def __init__(self, user_manager): + def __init__(self, user_manager, room_manager): self.user_manager = user_manager + self.room_manager = room_manager self.commands = { "help": self.cmd_help, "login": self.cmd_login, @@ -16,6 +17,9 @@ class CommandProcessor: "ban": self.cmd_ban, "broadcast": self.cmd_broadcast, "passwd": self.cmd_passwd, + "join": self.cmd_join, + "rooms": self.cmd_rooms, + "createroom": self.cmd_createroom, } def process_command(self, message: str, addr: Tuple) -> str: @@ -48,6 +52,8 @@ class CommandProcessor: auth_commands = { "broadcast": "Broadcast a message to all users", "users": "List online users", + "join": "Join a chat room", + "rooms": "List available rooms", } # Admin commands @@ -56,6 +62,7 @@ class CommandProcessor: "deop": "Remove admin privileges from user", "kick": "Disconnect a user from the server", "ban": "Ban a username from the server", + "createroom": "Create a new chat room", } # Build help message @@ -190,10 +197,9 @@ class CommandProcessor: if not args: return "Usage: broadcast " - # Check if user is authenticated (not a guest) - username = self.user_manager.get_username(addr) - if username.startswith("guest_"): - return "You must be logged in to broadcast messages" + # Check if user is admin + if not self.user_manager.is_admin(addr): + return "You must be an admin to broadcast messages" if self.user_manager.is_rate_limited(addr): return "Rate limit exceeded. Please wait a moment." @@ -217,3 +223,40 @@ class CommandProcessor: if self.user_manager.change_password(username, new_password): return "Password changed successfully" return "Failed to change password" + + def cmd_join(self, args, addr): + """Join a chat room.""" + if len(args) != 1: + return "Usage: join " + + room_name = args[0] + if self.room_manager.join_room(addr, room_name): + return f"Joined room: {room_name}" + return "Room not found" + + def cmd_rooms(self, args, addr): + """List available rooms.""" + rooms = self.room_manager.list_rooms() + current_room = self.room_manager.get_user_room(addr) + + room_list = ["+--------AVAILABLE ROOMS--------+"] + for name, desc in rooms.items(): + current = " (current)" if name == current_room else "" + room_list.append(f"# {name:<15} - {desc[:20]}{current}") + room_list.append("+-----------------------------+") + return "\n".join(room_list) + + def cmd_createroom(self, args, addr): + """Create a new room (admin only).""" + if not self.user_manager.is_admin(addr): + return "You don't have permission to create rooms" + + if len(args) < 1: + return "Usage: createroom [description]" + + name = args[0] + description = " ".join(args[1:]) if len(args) > 1 else "" + + if self.room_manager.create_room(name, description): + return f"Room {name} created successfully" + return "Room already exists" diff --git a/libs/room_manager.py b/libs/room_manager.py new file mode 100644 index 0000000..32453ac --- /dev/null +++ b/libs/room_manager.py @@ -0,0 +1,86 @@ +from pathlib import Path +import json + + +class Room: + def __init__(self, name, description=""): + self.name = name + self.description = description + self.users = set() # Set of addr tuples + + def add_user(self, addr): + self.users.add(addr) + + def remove_user(self, addr): + self.users.discard(addr) + + +class RoomManager: + def __init__(self): + self.rooms = {} # {name: Room} + self.user_rooms = {} # {addr: room_name} + self.rooms_file = Path("data/rooms.json") + self._load_rooms() + + def _load_rooms(self): + if self.rooms_file.exists(): + with open(self.rooms_file) as f: + rooms_data = json.load(f) + for name, data in rooms_data.items(): + self.rooms[name] = Room(name, data.get("description", "")) + else: + # Create default lounge + self.rooms_file.parent.mkdir(exist_ok=True) + self.rooms["lounge"] = Room("lounge", "The default chat room") + self._save_rooms() + + def _save_rooms(self): + rooms_data = { + name: {"description": room.description} for name, room in self.rooms.items() + } + with open(self.rooms_file, "w") as f: + json.dump(rooms_data, f, indent=2) + + def create_room(self, name, description=""): + if name.lower() in self.rooms: + return False + self.rooms[name.lower()] = Room(name, description) + self._save_rooms() + return True + + def join_room(self, addr, room_name): + """Join a chat room.""" + room_name = room_name.lower() + if room_name not in self.rooms: + return False + + # Get old room for notification + old_room = None + if addr in self.user_rooms: + old_room_name = self.user_rooms[addr] + old_room = self.rooms[old_room_name] + + # Remove from current room + self.leave_current_room(addr) + + # Add to new room + self.rooms[room_name].add_user(addr) + self.user_rooms[addr] = room_name + + return True + + def leave_current_room(self, addr): + if addr in self.user_rooms: + current_room = self.rooms[self.user_rooms[addr]] + current_room.remove_user(addr) + del self.user_rooms[addr] + + def get_room_users(self, room_name): + room = self.rooms.get(room_name.lower()) + return room.users if room else set() + + def get_user_room(self, addr): + return self.user_rooms.get(addr, "lounge") + + def list_rooms(self): + return {name: room.description for name, room in self.rooms.items()} diff --git a/main.py b/main.py index 3bee012..1857af3 100644 --- a/main.py +++ b/main.py @@ -6,6 +6,7 @@ from libs.broadcast import broadcast_message from libs.user_manager import UserManager from libs.process_message import CommandProcessor from libs.banner import load_banner +from libs.room_manager import RoomManager # Load environment variables load_dotenv() @@ -23,7 +24,8 @@ connections_lock = threading.Lock() # Thread-safe operations on active_connecti # Initialize user management user_manager = UserManager() -command_processor = CommandProcessor(user_manager) +room_manager = RoomManager() +command_processor = CommandProcessor(user_manager, room_manager) def handle_backspace(display_buffer, conn): @@ -104,7 +106,19 @@ def process_complete_line(line, addr, active_connections, conn): return # Broadcast the message - broadcast_message(active_connections, f"[{username}]: {message}") + current_room = room_manager.get_user_room(addr) + room = room_manager.rooms[current_room] + message_with_user = f"[{username}@{current_room}]: {message}" + + # Send to room members + broadcast_message( + active_connections, + message_with_user, + addr, + room, + ) + # Send back to sender + conn.sendall(f"\r\n{message_with_user}\r\n".encode("ascii")) except UnicodeDecodeError: print("UnicodeDecodeError: ", line) @@ -121,6 +135,7 @@ def handle_client(conn, addr): # Register as guest initially username = user_manager.register_session(addr) + room_manager.join_room(addr, "lounge") # Put user in default room # Check if username is banned if user_manager.is_banned(username): @@ -140,7 +155,10 @@ def handle_client(conn, addr): ) conn.sendall(welcome_msg.encode("ascii")) broadcast_message( - active_connections, f"* New user connected from {addr[0]} as {username}", addr + active_connections, + f"* New user connected from {addr[0]} as {username}", + addr, + system_msg=True, ) # Initialize buffers @@ -180,9 +198,13 @@ def handle_client(conn, addr): def cleanup_client_connection(addr): """Clean up resources when a client disconnects.""" print(f"Client {addr} disconnected") + username = user_manager.get_username(addr) + room_manager.leave_current_room(addr) with connections_lock: del active_connections[addr] - broadcast_message(active_connections, f"* User {addr[0]} disconnected") + broadcast_message( + active_connections, f"* User {username} disconnected", system_msg=True + ) def start_server():