server.py

import threading
import socket
import argparse
import os


class Server(threading.Thread):
    def __init__(self, host, port):
        super().__init__()
        self.connections = []
        self.host = host
        self.port = port
        self.sock = None
        self.connections_lock = threading.Lock()

    def run(self):
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.sock.bind((self.host, self.port))
            self.sock.listen(1)
            print("Listening at", self.sock.getsockname())
        except OSError as e:
            print(f"Error starting server: {e}")
            self.shutdown()
            return

        while True:
            try:
                sc, sockname = self.sock.accept()
                print(f"Accepted a new connection from {sc.getpeername()} to {sc.getsockname()}")

                server_socket = ServerSocket(sc, sockname, self)
                server_socket.start()
                with self.connections_lock:
                    self.connections.append(server_socket)
                print("Ready to receive messages from", sc.getpeername())
            except OSError:
                # This will be raised when the socket is closed by the shutdown method
                break
        print("Server has been shut down.")

    def broadcast(self, message, source):
        with self.connections_lock:
            for connection in self.connections:
                if connection.sockname != source:
                    connection.send(message)

    def remove_connection(self, connection):
        with self.connections_lock:
            if connection in self.connections:
                self.connections.remove(connection)

    def shutdown(self):
        print("Closing all connections...")
        with self.connections_lock:
            for connection in self.connections:
                connection.sc.close()
        print("Shutting down the server...")
        if self.sock:
            self.sock.close()


class ServerSocket(threading.Thread):
    def __init__(self, sc, sockname, server):
        super().__init__()
        self.sc = sc
        self.sockname = sockname
        self.server = server

    def run(self):
        try:
            while True:
                message = self.sc.recv(1024).decode("utf-8")
                if message:
                    print(f"{self.sockname} says {message}")
                    self.server.broadcast(message, self.sockname)
                else:
                    print(f"Connection gracefully closed by {self.sockname}")
                    break
        except (ConnectionResetError, BrokenPipeError, UnicodeDecodeError) as e:
            print(f"Connection lost with {self.sockname}: {e}")
        finally:
            self.sc.close()
            self.server.remove_connection(self)

    def send(self, message):
        self.sc.sendall(message.encode("utf-8"))


def handle_shutdown(server):
    while True:
        ipt = input("")
        if ipt == "q":
            server.shutdown() 
            break


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Chatroom Server")
    parser.add_argument("host", help="Interface the server listens at")
    parser.add_argument('-p', metavar='port', type=int, default=1060, help="TCP port (default 1060)")

    args = parser.parse_args()

    server = Server(args.host, args.p)
    server.start()

    shutdown_thread = threading.Thread(target=handle_shutdown, args=(server,))
    shutdown_thread.daemon = True  # Allows main program to exit even if this thread is running
    shutdown_thread.start()