server.py

import socket
import threading
import ssl
import time
from protocol import send_json, recv_json
from encryption_utils import generate_key

HOST = "0.0.0.0"
PORT = 5000

clients = {}
user_sockets = {}
rooms = {}
lock = threading.Lock()

stats = {
    "messages": 0,
    "files": 0,
    "connections": 0
}

def broadcast_room(room, message, exclude=None):
    if room in rooms:
        for client in rooms[room]:
            if client != exclude:
                try:
                    send_json(client, message)
                except:
                    remove_client(client)

def remove_client(client):
    with lock:
        username = clients.get(client)
        if username:
            print(f"{username} disconnected")
            user_sockets.pop(username, None)
            clients.pop(client, None)

            for room in rooms.values():
                if client in room:
                    room.remove(client)

    client.close()

def handle_client(client):
    buffer = ""

    while True:
        try:
            messages, buffer = recv_json(client, buffer)

            for data in messages:
                msg_type = data.get("type")

                if msg_type == "join":
                    username = data["username"]
                    room = data["room"]

                    with lock:
                        clients[client] = username
                        user_sockets[username] = client

                        if room not in rooms:
                            rooms[room] = []
                        rooms[room].append(client)

                    send_json(client, {"type": "info", "message": f"Joined {room}"})
                    broadcast_room(room, {"type": "info", "message": f"{username} joined"}, client)

                elif msg_type == "message":
                    stats["messages"] += 1
                    room = data["room"]
                    username = clients.get(client)

                    broadcast_room(room, {
                        "type": "message",
                        "sender": username,
                        "content": data["content"]
                    })

                elif msg_type == "private":
                    to_user = data["to"]
                    if to_user in user_sockets:
                        send_json(user_sockets[to_user], {
                            "type": "private",
                            "from": clients[client],
                            "content": data["content"]
                        })

                elif msg_type == "secret":
                    to_user = data["to"]
                    if to_user in user_sockets:
                        send_json(user_sockets[to_user], {
                            "type": "secret",
                            "from": clients[client],
                            "content": data["content"]
                        })

                elif msg_type == "file":
                    stats["files"] += 1

                    filename = data["filename"]
                    filesize = data["filesize"]
                    room = data["room"]
                    username = clients.get(client)

                    broadcast_room(room, {
                        "type": "file",
                        "filename": filename,
                        "filesize": filesize,
                        "from": username
                    }, client)

                    remaining = filesize
                    while remaining > 0:
                        chunk = client.recv(min(4096, remaining))
                        remaining -= len(chunk)

                        for member in rooms[room]:
                            if member != client:
                                try:
                                    member.sendall(chunk)
                                except:
                                    remove_client(member)

        except Exception as e:
            print("Client error:", e)
            break

    remove_client(client)

def monitor():
    while True:
        print(f"[STATS] Clients: {len(clients)} | Messages: {stats['messages']} | Files: {stats['files']}")
        time.sleep(10)

def start_server():
    generate_key()

    context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    context.load_cert_chain(certfile="cert.pem", keyfile="key.pem")

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind((HOST, PORT))
    server.listen(5)

    print("Server running with SSL...")

    threading.Thread(target=monitor, daemon=True).start()

    while True:
        client, addr = server.accept()
        secure_client = context.wrap_socket(client, server_side=True)

        print(f"Connection from {addr} | Cipher: {secure_client.cipher()}")
        stats["connections"] += 1

        threading.Thread(target=handle_client, args=(secure_client,), daemon=True).start()

if __name__ == "__main__":
    start_server()