server.py

import socket
import threading
import time
import hamming


class Server:
    """
    A multi-threaded TCP chat server that handles client connections,
    message framing, and command processing.

    Here's the basic flow of how the server works:

    1. The start() method ... starts things up. It binds the server to a
       host and port and enters a permanent loop to listen for new client
       connections.

    2. When a new client connects, the server hands them off to a
       new, dedicated thread running the _handle_client method. This lets
       the main thread go back to listening for more connections.

    3. The _handle_client method's job is to listen for data from its one
       specific client. It uses a buffer to piece together complete messages,
       which are expected to be framed by | characters.

    4. Once a full message is extracted, it's passed to _process_payload.
       This function acts as a router. If the client isn't registered yet,
       it expects to see an IAM command as the first message, which is handled by _handle_iam. 
       Otherwise, it's disconnected.

    5. If the client is already registered, _process_payload checks the
       command (e.g., "to", "groupsend") and uses the self.handlers
       dictionary to call the appropriate function to deal with the request.
       Otherwise it calls _handle_unknown to let the client know they messed up.
    """


    def __init__(self, host="0.0.0.0", port=12345):
        self.host = host
        self.port = port
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        # Clients are stored as :
        # {conn: {"name": str, "addr": tuple, "pk_e": int, "pk_n": int}}
        self.clients = {}

        self.lock = threading.Lock()

        # Command dispatcher for registered clients
        self.handlers = {
            "to": self._handle_to,
            "groupsend": self._handle_groupsend,
        }

    def log(self, message):
        """Prints a timestamped log message."""

        timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        print(f"[{timestamp}] {message}")

    def broadcast(self, message, exclude_conn=None):
        """Sends a message to all connected clients, optionally excluding one."""

        with self.lock:
            # Iterate over a copy of items to avoid issues if dict changes
            for conn, client_data in list(self.clients.items()):
                if conn != exclude_conn:
                    try:
                        conn.sendall(f"|{message}|".encode())
                    except Exception as e:
                        self.log(
                            f"Error broadcasting to {client_data.get('name', 'Unknown')}: {e}"
                        )

    def start(self):
        """Binds the server and starts the main loop to accept connections."""

        try:
            self.server_socket.bind((self.host, self.port))
            self.server_socket.listen(5)
            self.log(f"Server listening on {self.host}:{self.port}")

            while True:
                conn, addr = self.server_socket.accept()

                # Hand off each client to a separate thread
                thread = threading.Thread(
                    target=self._handle_client, args=(conn, addr), daemon=True
                )
                thread.start()

        except OSError as e:
            self.log(f"Failed to bind server to {self.host}:{self.port} - {e}")
        except KeyboardInterrupt:
            self.log("Admin shutdown initiated...")
        finally:
            self.shutdown()

    def shutdown(self):
        """Gracefully shuts down the server."""
        self.log("Server is shutting down.")
        self.broadcast("system:shutdown:Server is shutting down.")
        time.sleep(0.5)  # Give clients a moment to receive the message

        with self.lock:
            for conn in list(self.clients.keys()):
                conn.close()
            self.clients.clear()

        self.server_socket.close()
        self.log("Server has shut down completely.")

    def _remove_client(self, conn):
        """Removes a client from the list and notifies others."""

        with self.lock:
            client_data = self.clients.pop(conn, None)

        if client_data:
            self.log(f"Connection closed for {client_data['name']}")
            self.broadcast(f"system:left:{client_data['name']}")
        else:
            self.log(f"Connection closed for a non-registered client.")

        try:
            conn.close()
        except Exception:
            pass

    def _handle_client(self, conn, addr):
        """Manages a single client connection, processing incoming messages."""
        self.log(f"New connection from {addr}")

        # Maintain a buffer to handle messages that may arrive across multiple recv() calls
        # After each recv(), attempt to extract a complete '|' delimited message from the buffer
        # If successful, remove it from the buffer and process it
        receive_buffer = ""
        try:
            while True:
                data = conn.recv(4096)
                if not data:
                    break

                receive_buffer += data.decode()
                while "|" in receive_buffer:
                    try:
                        start = receive_buffer.index("|")
                        end = receive_buffer.index("|", start + 1)
                        payload = receive_buffer[start + 1 : end]
                        receive_buffer = receive_buffer[end + 1 :]
                        self._process_payload(conn, addr, payload)
                    except ValueError:
                        break  # Incomplete message

        except ConnectionResetError:
            self.log(f"Connection reset by client at {addr}")
        except Exception as e:
            self.log(f"Unhandled error with client {addr}: {e}")
        finally:
            self._remove_client(conn)

    def _process_payload(self, conn, addr, payload):
        """Parses a payload and routes it to the appropriate handler."""

        parts = payload.split(":", 3)
        command = parts[0]

        with self.lock:
            client_data = self.clients.get(conn)

        if client_data is None:  # Client not yet registered
            # Expect a IAM identification message
            if command == "IAM":
                self._handle_iam(conn, addr, parts)
            else:
                # Client has goofed
                self.log(f"Invalid initial message from {addr}: {payload}. Closing.")
                conn.sendall(b"|system:error:Invalid initial message format|")
                self._remove_client(conn)  # Force disconnect

        else:  # Client is registered
            handler = self.handlers.get(command, self._handle_unknown)
            handler(client_data, parts)

    # --- Command Handler Methods ---

    def _handle_iam(self, conn, addr, parts):
        """Handles the initial client registration."""

        if len(parts) != 4:
            conn.sendall(b"|system:error:Malformed IAM message|")
            return

        name, pk_e_str, pk_n_str = parts[1], parts[2], parts[3]
        with self.lock:
            if any(c["name"] == name for c in self.clients.values()):
                self.log(f"Username '{name}' taken. Rejecting {addr}.")
                conn.sendall(b"|username_taken|")
                return

            self.clients[conn] = {
                "name": name,
                "addr": addr,
                "pk_e": int(pk_e_str),
                "pk_n": int(pk_n_str),
                "conn": conn,
            }

            # Prepare list of other users and their key pairs
            # to send to the new client
            users_list = [
                f"{c['name']}/{c['pk_e']}/{c['pk_n']}"
                for c in self.clients.values()
                if c["conn"] != conn
            ]
            users_str = ",".join(users_list)

        welcome_msg = f"welcome:Welcome {name}!\nusers:{users_str}"
        conn.sendall(f"|{welcome_msg}|".encode())
        self.log(f"User {name} from {addr} registered.")

        # Let other clients know the key pair of the new guy
        broadcast_msg = f"system:joined:{name}:{pk_e_str}:{pk_n_str}"
        self.broadcast(broadcast_msg, exclude_conn=conn)

    def _handle_to(self, sender_data, parts):
        """Handles a direct message from one client to another."""

        if len(parts) < 3:
            return  # malformed

        recipient_name, encrypted_msg = parts[1], parts[2]

        with self.lock:
            recipient = next(
                (c for c in self.clients.values() if c["name"] == recipient_name), None
            )

        # If recepient exists, simply relay the message else inform sender
        if recipient and recipient.get("conn"):
            try:
                msg = f"|from:{sender_data['name']}:{encrypted_msg}|"
                recipient["conn"].sendall(msg.encode())
                self.log(
                    f"Relayed message from {sender_data['name']} to {recipient_name}"
                )
            except Exception as e:
                self.log(f"Error sending to {recipient_name}: {e}")
                sender_data["conn"].sendall(
                    f"|system:error:User {recipient_name} disconnected.|".encode()
                )
        else:
            self.log(
                f"User {recipient_name} not found for message from {sender_data['name']}."
            )
            sender_data["conn"].sendall(
                f"|system:error:User {recipient_name} not online.|".encode()
            )

    def _handle_groupsend(self, sender_data, parts):
        """Handles a group message, flips a bit, and broadcasts."""

        if len(parts) < 3:
            return  # malformed

        data_len, group_message = parts[1], parts[2]

        # Simulate nOiSy cHaNnEl by flipping a *single* bit in the entire message
        group_message, indices = hamming.flip_bits(group_message, 1)
        print(f"Flipping bit : {indices[0]} in message from {sender_data['name']}")

        self.log(f"Broadcasting group message from {sender_data['name']}")
        self.broadcast(
            f"groupmsg:{sender_data['name']}:{data_len}:{group_message}",
            exclude_conn=sender_data["conn"],
        )

    def _handle_unknown(self, sender_data, parts):
        """Handles any unrecognized command from a registered client."""

        command = parts[0]
        self.log(f"Unknown command '{command}' from {sender_data['name']}")
        sender_data["conn"].sendall(
            f"|system:error:Unknown command: {command}|".encode()
        )


if __name__ == "__main__":
    server = Server()
    server.start()