cw-generator/signalsrv/signalsrv.py

import asyncio
import datetime
import json
import logging
import re

from websockets.asyncio.server import serve

__VERSION__ = "0.0.1"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(name)s %(levelname)s %(message)s"
)


class Client:
    freqs = {}
    freq_re = re.compile(r"^\d+\.\d{3}$")

    def __init__(self, websocket):
        self.websocket = websocket
        self.curr_freq = None

    async def handle(self):
        print(f" >>> New client {self.client} connected")
        exc = None
        try:
            await self._handle_client()
        except Exception as e:
            exc = e
        finally:
            # FIXME: basically handle disconnect / leave from room
            print(f" <<< Client {self.client} id {self.id} disconnected: {exc}")
            if self.curr_freq:
                await self._leave_room()

    @property
    def client(self):
        ip, port, *_ = self.websocket.remote_address
        if ':' in ip:
            ip = f"[{ip}]"
        return f"{ip}:{port}"

    @property
    def id(self):
        return str(self.websocket.id)

    async def _handle_client(self):
        await self._send(type="hello", name="LobbySrv 3000", version=__VERSION__)
        async for data in self.websocket:
            print(f" <-- client {self.client} sent {repr(data)}")
            try:
                data = json.loads(data)
            except json.JSONDecodeError:
                self._send_error("Could not decode message, invalid json")
                continue

            if not isinstance(data, dict) or "cmd" not in data:
                await self._send_error("Invalid format in json")
                continue

            print(f"{datetime.datetime.now()} {self.client} wrote:", data)

            match data["cmd"]:
                case "quit":
                    break
                case "create":
                    await self._create_room(data)
                case "join":
                    await self._join_room(data)
                case "leave":
                    await self._leave_room()
                case "list":
                    freqs = [{"freq": freq, "players": len(players)}
                             for freq, players in self.freqs.items()]
                    await self._send(type="freq-list", freqs=freqs)
                case "disconnect":
                    pass
                case "morse-state":
                    await self._handle_morse_state(data)
                case _:
                    await self._send_error("Unknown command")

    async def _create_room(self, data):
        if self.curr_freq:
            await self._send_error(f"Already on frequency {self.curr_freq}")
            return

        if "freq" not in data:
            await self._send_error("No frequency in create message")
            return
        freq = data["freq"]

        if not self.freq_re.match(freq):
            await self._send_error("Invalid frequency")
            return

        if freq in self.freqs:
            await self._send_error("Frequency already in use")
            return

        self.curr_freq = freq
        self.freqs[freq] = [self]
        await self._send(type="join", freq=self.curr_freq, self_id=self.id, other_players=[])

    async def _join_room(self, data):
        if self.curr_freq:
            await self._send_error(f"Already on frequency {self.curr_freq}")
            return

        if "freq" not in data:
            await self._send_error("No frequency in join message")
            return
        freq = data["freq"]

        if freq not in self.freqs:
            await self._send_error(f"Frequency {freq} not available")
            return

        self.curr_freq = freq
        self.freqs[freq].append(self)
        # FIXME: do we need locking here?
        print("FREQ", self.curr_freq, freq, self.freqs)
        await self._send(type="join", freq=self.curr_freq, self_id=self.id,
                         other_players=[c.id for c in self._others(freq)])
        await self._send_to_group(self._others(freq), type="player-joined", player=self.id)

    async def _handle_morse_state(self, data):
        if not self.curr_freq:
            await self._send_error("No frequency selected")
            return

        if "state" not in data or not isinstance(data["state"], bool):
            await self._send_error("No state key with type bool in data")
            return

        await self._send_to_group(self._others(self.curr_freq),
                                  type="morse-state", state=data["state"], from_player=self.id)

    async def _leave_room(self):
        if not self.curr_freq:
            self._send_error("You are not on a frequency")
            return

        await self._send_to_group(self._others(self.curr_freq),
                                  type="player-left", player=self.id)
        try:
            self.freqs[self.curr_freq].remove(self)
        except ValueError:
            print(f"Warning: Player {self.id} was not in freq {self.curr_freq}")
        if not self.freqs[self.curr_freq]:
            del self.freqs[self.curr_freq]
        self.curr_freq = None

        try:
            await self._send(type="leave")
        except Exception:
            pass

    def _others(self, freq):
        return [c for c in self.freqs[freq] if c.id != self.id]

    async def _send(self, ignore_exceptions=False, **kwargs):
        data = json.dumps(kwargs).encode()
        print(f" --> sending out to {self.client}: {data}")
        try:
            await self.websocket.send(json.dumps(kwargs).encode() + b"\n")
        except Exception as e:
            print(f"Error sending data to {self.client}: {e}")
            if not ignore_exceptions:
                raise

    async def _send_to_group(self, group, **kwargs):
        async with asyncio.TaskGroup() as tg:
            for member in group:
                tg.create_task(member._send(ignore_exceptions=True, **kwargs))

    async def _send_error(self, msg: str):
        await self._send(type="error", message=msg)


async def new_client(websocket):
    try:
        client = Client(websocket)
        await client.handle()
    finally:
        pass
    # async for message in websocket:
    #    await websocket.send(message)


async def main():
    HOST, PORT = "0.0.0.0", 3784
    async with serve(new_client, HOST, PORT) as server:
        await server.serve_forever()


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