server.py

#!/usr/bin/env python3
"""
keycoms — server.py
WebSocket server: manages rooms, roles, and keystroke broadcasting.
"""

import asyncio
import json
import uuid
import websockets
from websockets.server import WebSocketServerProtocol

# rooms[room_code] = {
#   "admin": ws,
#   "users": { ws: { "id": str, "role": "streamer"|"listener" } }
# }
rooms: dict = {}


def generate_room_code() -> str:
    return uuid.uuid4().hex[:6].upper()


def get_user(room_code: str, ws: WebSocketServerProtocol) -> dict | None:
    room = rooms.get(room_code)
    if room:
        return room["users"].get(ws)
    return None


def find_ws_by_id(room_code: str, user_id: str) -> WebSocketServerProtocol | None:
    room = rooms.get(room_code)
    if not room:
        return None
    for ws, info in room["users"].items():
        if info["id"] == user_id:
            return ws
    return None


async def broadcast(room_code: str, message: dict, exclude: WebSocketServerProtocol = None):
    room = rooms.get(room_code)
    if not room:
        return
    payload = json.dumps(message)
    targets = [ws for ws in room["users"] if ws != exclude]
    if targets:
        await asyncio.gather(*[ws.send(payload) for ws in targets], return_exceptions=True)


async def send(ws: WebSocketServerProtocol, message: dict):
    await ws.send(json.dumps(message))


async def handle_command(ws: WebSocketServerProtocol, room_code: str, data: dict):
    room = rooms[room_code]
    is_admin = room["admin"] == ws
    action = data.get("action")
    target_id = data.get("target")

    if not is_admin:
        await send(ws, {"type": "error", "msg": "Not authorized."})
        return

    if action == "list":
        users = [
            {"id": info["id"], "role": info["role"]}
            for info in room["users"].values()
        ]
        await send(ws, {"type": "list", "users": users})

    elif action in ("mute", "unmute"):
        target_ws = find_ws_by_id(room_code, target_id)
        if not target_ws:
            await send(ws, {"type": "error", "msg": f"User {target_id} not found."})
            return
        new_role = "listener" if action == "mute" else "streamer"
        room["users"][target_ws]["role"] = new_role
        await send(target_ws, {"type": "role", "role": new_role})
        await send(ws, {"type": "info", "msg": f"{target_id} is now a {new_role}."})

    elif action == "kick":
        target_ws = find_ws_by_id(room_code, target_id)
        if not target_ws:
            await send(ws, {"type": "error", "msg": f"User {target_id} not found."})
            return
        await send(target_ws, {"type": "kicked", "msg": "You have been kicked."})
        await target_ws.close()

    else:
        await send(ws, {"type": "error", "msg": f"Unknown action: {action}"})


async def handler(ws: WebSocketServerProtocol):
    user_id = uuid.uuid4().hex[:8]
    room_code = None

    try:
        # Step 1: expect join message
        raw = await ws.recv()
        data = json.loads(raw)

        if data.get("type") != "join":
            await send(ws, {"type": "error", "msg": "Expected join message."})
            return

        room_code = data.get("room", "").strip().upper()
        create_new = data.get("create", False)

        if create_new or room_code not in rooms:
            # Create room
            if not room_code:
                room_code = generate_room_code()
            rooms[room_code] = {
                "admin": ws,
                "users": {}
            }
            role = "streamer"
            is_admin = True
        else:
            role = "listener"
            is_admin = False

        rooms[room_code]["users"][ws] = {"id": user_id, "role": role}

        await send(ws, {
            "type": "joined",
            "room": room_code,
            "id": user_id,
            "role": role,
            "admin": is_admin
        })

        print(f"[+] {user_id} joined room {room_code} as {'ADMIN/' if is_admin else ''}{role}")

        # Notify others
        await broadcast(room_code, {
            "type": "peer_joined",
            "id": user_id,
            "role": role
        }, exclude=ws)

        # Step 2: message loop
        async for raw in ws:
            try:
                data = json.loads(raw)
            except json.JSONDecodeError:
                continue

            msg_type = data.get("type")

            if msg_type == "key":
                user = get_user(room_code, ws)
                if user and user["role"] == "streamer":
                    # Broadcast to all listeners
                    room = rooms[room_code]
                    payload = json.dumps({
                        "type": "key",
                        "value": data.get("value"),
                        "from": user_id
                    })
                    targets = [
                        w for w, info in room["users"].items()
                        if info["role"] == "listener" and w != ws
                    ]
                    if targets:
                        await asyncio.gather(
                            *[w.send(payload) for w in targets],
                            return_exceptions=True
                        )

            elif msg_type == "cmd":
                await handle_command(ws, room_code, data)

    except websockets.exceptions.ConnectionClosed:
        pass
    finally:
        if room_code and room_code in rooms:
            room = rooms[room_code]
            room["users"].pop(ws, None)
            print(f"[-] {user_id} left room {room_code}")

            # Notify others
            await broadcast(room_code, {"type": "peer_left", "id": user_id})

            # If admin left, assign new admin or clean room
            if room["admin"] == ws:
                if room["users"]:
                    new_admin_ws = next(iter(room["users"]))
                    room["admin"] = new_admin_ws
                    room["users"][new_admin_ws]["role"] = "streamer"
                    await send(new_admin_ws, {
                        "type": "promoted",
                        "msg": "You are now the admin."
                    })
                    print(f"[*] {room['users'][new_admin_ws]['id']} promoted to admin in {room_code}")
                else:
                    del rooms[room_code]
                    print(f"[*] Room {room_code} dissolved.")


async def main():
    host = "0.0.0.0"
    port = 8765
    print(f"[keycoms] Server running on ws://{host}:{port}")
    print(f"[keycoms] Tunnel with: cloudflared tunnel --url ws://localhost:{port}")
    print()
    async with websockets.serve(handler, host, port):
        await asyncio.Future()


if __name__ == "__main__":
    asyncio.run(main())