338class Game:
339    """Subclass to implement game logic.
340
341    This class should not be initialized directly, it is initialized by the server.
342    """
343
344    persistent: bool = False
345    """Set as persistent to allow the game to persist even without players.
346
347    This is required for saving and loading (see also: `Game.get_save_string`).
348    """
349    heartbeat_rate: float = 10
350    """How many times per second the client should check for updates.
351
352    See `Game.handle_heartbeat`.
353    """
354
355    def __init__(self, name: str, save_string: Optional[str] = None):
356        """The server initializes this class for every game started.
357
358        Args:
359            name: Game instance name.
360            save_string: Game data loaded from disk from last server session as given by
361                `Game.get_save_string`.
362        """
363        pass
364
365    def user_joined(self, username: str):
366        """Called when a user joins the game."""
367        pass
368
369    def user_left(self, username: str):
370        """Called when a user leaves the game."""
371        pass
372
373    def handle_packet(self, packet: Packet) -> Response:
374        """Packet handling for heartbeat updates and game requests.
375
376        Most use cases should override `Game.handle_game_packet` and
377        `Game.handle_heartbeat` instead of this method.
378        """
379        if packet.message == Request.HEARTBEAT_UPDATE:
380            return self.handle_heartbeat(packet)
381        return self.handle_game_packet(packet)
382
383    def handle_game_packet(self, packet: Packet) -> Response:
384        """Override this method to implement packet handling.
385
386        See also: `pgnet.Client.send`.
387        """
388        return Response(
389            f"No packet handling configured for {self.__class__.__qualname__}",
390            status=Status.UNEXPECTED,
391        )
392
393    def handle_heartbeat(self, packet: Packet) -> Response:
394        """Override this method to implement heartbeat updates.
395
396        See also: `Game.heartbeat_rate`, `pgnet.Client.heartbeat_payload`,
397        `pgnet.Client.on_heartbeat`.
398        """
399        return Response(
400            f"No heartbeat update configured for {self.__class__.__qualname__}"
401        )
402
403    def get_save_string(self) -> str:
404        """Override this method to save game data to disk.
405
406        If `Game.persistent`, this method is called by the server periodically and when
407        shutting down. In the next session, the server will recreate the game with this
408        string passed as *`save_string`* to `Game.__init__`.
409
410        .. note:: `pgnet.Server` must be configured to enable saving and loading.
411        """
412        return ""
413
414    def get_lobby_info(self) -> str:
415        """Override this method to post public game info in lobby.
416
417        This will be used to give users more information about a game before joining.
418        """
419        return ""
420
421    def update(self):
422        """Called on an interval by the server.
423
424        Override this method to implement background game logic tasks.
425        """
426        pass

 53class Client:
 54    """The client that manages communication with the server.
 55
 56    ## Initializing
 57    This class should not be initialized directly, instead use `Client.remote` or
 58    `Client.local`.
 59    ```python3
 60    # This client will connect to a remote server
 61    remote_client = Client.remote(address="1.2.3.4", username="player")
 62    # This client will create a local server and connect to it
 63    local_client = Client.local(game=MyGame, username="player")
 64    ```
 65
 66    ## Connecting and starting a game
 67    Use `Client.async_connect` to connect to the server. Once connected, use
 68    `Client.get_game_dir`, `Client.create_game`, `Client.join_game`, and
 69    `Client.leave_game` to join or leave a game.
 70
 71    ## Using the packet queue
 72    When in game (`Client.game` is not *None*), use `Client.send` to send a
 73    `pgnet.Packet` to `pgnet.Game.handle_game_packet` and receive a `pgnet.Response`.
 74
 75    ## Client events
 76    It is possible to bind callbacks to client events by subclassing and overriding or
 77    by setting `Client.on_connection`, `Client.on_status`, and `Client.on_game`.
 78
 79    """
 80
 81    def __init__(
 82        self,
 83        *,
 84        address: str,
 85        username: str,
 86        password: str,
 87        invite_code: str,
 88        port: int,
 89        server: Optional[Server],
 90        verify_server_pubkey: str,
 91    ):
 92        """This class should not be initialized directly.
 93
 94        .. note:: Use `Client.local` or `Client.remote` to create a client.
 95        """
 96        self._key: Key = Key()
 97        self._status: str = "New client."
 98        self._server_connection: Optional[WebSocketClientProtocol] = None
 99        self._connected: bool = False
100        self._do_disconnect: bool = False
101        self._packet_queue: list[tuple[Packet, ResponseCallback]] = []
102        self._game: Optional[str] = None
103        self._heartbeat_interval = 1 / DEFAULT_HEARTBEAT_RATE
104        # Connection details
105        self._username = username
106        self._password = password
107        self._invite_code = invite_code
108        self._address = address
109        self._port = port
110        self._server = server
111        self._verify_server_pubkey = verify_server_pubkey
112
113    @classmethod
114    def local(
115        cls,
116        *,
117        game: Type[Game],
118        username: str,
119        password: str = "",
120        port: int = DEFAULT_PORT,
121        server_factory: Optional[Callable[[Any], Server]] = None,
122    ) -> "Client":
123        """Create a client that uses its own local server. See also `Client.remote`.
124
125        Only the *game* and *username* arguments are required.
126
127        Args:
128            game: `pgnet.Game` class to pass to the local server.
129            username: The user's username.
130            password: The user's password.
131            port: Server port number.
132            server_factory: If provided, will be used to create the local server. Must
133                accept the same arguments as `pgnet.Server` (default). This is useful
134                for using a server subclass or to pass custom arguments to the local
135                server.
136        """
137        server_factory = server_factory or Server
138        server = server_factory(
139            game,
140            listen_globally=False,
141            registration_enabled=True,
142            require_user_password=False,
143        )
144        return cls(
145            address="localhost",
146            username=username,
147            password=password,
148            invite_code="",
149            port=port,
150            server=server,
151            verify_server_pubkey=server.pubkey,
152        )
153
154    @classmethod
155    def remote(
156        cls,
157        *,
158        address: str,
159        username: str,
160        password: str = "",
161        invite_code: str = "",
162        port: int = DEFAULT_PORT,
163        verify_server_pubkey: str = "",
164    ) -> "Client":
165        """Create a client that connects to a remote server. See also `Client.local`.
166
167        Args:
168            address: Server IP address.
169            username: The user's username.
170            password: The user's password.
171            invite_code: An invite code for creating a new user.
172            port: Server port number.
173            verify_server_pubkey: If provided, will compare against the
174                `pgnet.Server.pubkey` of the server and disconnect if they do not match.
175        """
176        return cls(
177            address=address,
178            username=username,
179            password=password,
180            invite_code=invite_code,
181            port=port,
182            server=None,
183            verify_server_pubkey=verify_server_pubkey,
184        )
185
186    async def async_connect(self):
187        """Connect to a server.
188
189        This procedure will automatically create an end-to-end encrypted connection
190        (optionally verifying the public key first), and authenticate the username and
191        password. Only after succesfully completing these steps will the client be
192        considered connected and ready to process the packet queue from `Client.send`.
193        """
194        logger.debug(f"Connecting... {self}")
195        if self._server is None:
196            return await self._async_connect_remote()
197        else:
198            return await self._async_connect_local()
199
200    def disconnect(self):
201        """Close the connection."""
202        self._do_disconnect = True
203
204    def get_game_dir(self, callback: Callable, /):
205        """Get the games directory from the server and pass the response to callback."""
206        self.send(Packet(Request.GAME_DIR), callback, do_next=True)
207
208    def create_game(
209        self,
210        name: str,
211        /,
212        *,
213        password: Optional[str] = None,
214        callback: Optional[ResponseCallback] = None,
215    ):
216        """Request from the server to create and join a game."""
217        payload = dict(name=name)
218        if password:
219            payload["password"] = password
220        self.send(Packet(Request.CREATE_GAME, payload), callback, do_next=True)
221
222    def join_game(
223        self,
224        name: str,
225        /,
226        *,
227        password: Optional[str] = None,
228        callback: Optional[ResponseCallback] = None,
229    ):
230        """Request from the server to join a game."""
231        payload = dict(name=name)
232        if password:
233            payload["password"] = password
234        self.send(Packet(Request.JOIN_GAME, payload), callback, do_next=True)
235
236    def leave_game(
237        self,
238        *,
239        callback: Optional[ResponseCallback] = None,
240    ):
241        """Request from the server to leave the game."""
242        self.send(Packet(Request.LEAVE_GAME), callback, do_next=True)
243
244    def send(
245        self,
246        packet: Packet,
247        callback: Optional[ResponseCallback] = None,
248        /,
249        do_next: bool = False,
250    ):
251        """Add a `pgnet.Packet` to the queue.
252
253        If a callback was given, the `pgnet.Response` will be passed to it. Callbacks
254        are not ensured, as the queue can be arbitrarily cleared using
255        `Client.flush_queue`.
256        """
257        if not self._connected:
258            logger.warning(f"Cannot queue packets while disconnected: {packet}")
259            return
260        packet_count = len(self._packet_queue)
261        if packet_count >= 5:
262            logger.warning(f"{packet_count} packets pending, adding: {packet}")
263        packet_callback = (packet, callback)
264        if do_next:
265            self._packet_queue.insert(0, packet_callback)
266        else:
267            self._packet_queue.append(packet_callback)
268
269    def flush_queue(self):
270        """Clear packets and their respective callbacks from the queue.
271
272        See also: `Client.send`.
273        """
274        if self._packet_queue:
275            logger.debug(f"Discarding messages:  {self._packet_queue}")
276        self._packet_queue = []
277
278    def on_connection(self, connected: bool):
279        """Called when connected or disconnected. See also: `Client.connected`."""
280        pass
281
282    @property
283    def connected(self) -> bool:
284        """If we are connected to the server. See also: `Client.on_connection`."""
285        return self._connected
286
287    def on_status(self, status: str):
288        """Called with feedback on client status. See also: `Client.status`."""
289        pass
290
291    @property
292    def status(self) -> str:
293        """Last status feedback message. See also: `Client.on_status`."""
294        return self._status
295
296    def on_game(self, game_name: Optional[str]):
297        """Called when joining or leaving a game. See also: `Client.game`."""
298        pass
299
300    @property
301    def game(self) -> Optional[str]:
302        """Currently joined game name. See also: `Client.on_game`."""
303        return self._game
304
305    def on_heartbeat(self, heartbeat: Response):
306        """Override this method to implement heartbeat updates.
307
308        The *heartbeat* Response is given by `pgnet.util.Game.handle_heartbeat`. The
309        heartbeat rate is set by `pgnet.util.Game.heartbeat_rate`.
310
311        See also: `Client.heartbeat_payload`.
312        """
313        pass
314
315    def heartbeat_payload(self) -> dict:
316        """Override this method to add data to the heartbeat request payload.
317
318        This payload is passed to `pgnet.util.Game.handle_heartbeat`. The heartbeat rate
319        is set by `pgnet.util.Game.heartbeat_rate`.
320
321        See also: `Client.on_heartbeat`.
322        """
323        return dict()
324
325    async def _async_connect_remote(self):
326        """Connect to the server.
327
328        Gets a websocket to create a `ClientConnection` object. This is passed to the
329        handshake handler to be populated, and then to the user connection for handling
330        the packet queue. The heartbeat task is managed here.
331        """
332        if self._server_connection is not None:
333            raise RuntimeError("Cannot open more than one connection per client.")
334        logger.debug("Connecting to server...")
335        full_address = f"{self._address}:{self._port}"
336        self._server_connection = websockets.connect(
337            f"ws://{full_address}",
338            close_timeout=1,
339        )
340        self._set_status(f"Connecting to {full_address}...", logger.info)
341        connection: Optional[ClientConnection] = None
342        heartbeat: Optional[asyncio.Task] = None
343        try:
344            try:
345                websocket = await self._server_connection
346            except OSError as e:
347                logger.debug(f"{e=}")
348                raise DisconnectedError("Failed to call server.")
349            connection = ClientConnection(websocket)
350            self._set_status("Logging in...", logger.info)
351            await self._handle_handshake(connection)
352            self._set_connection(True)
353            self._set_status(
354                f"Logged in as {self._username} @ {connection.remote}",
355                logger.info,
356            )
357            heartbeat = asyncio.create_task(self._async_heartbeat())
358            await self._handle_user_connection(connection)
359        except DisconnectedError as e:
360            logger.debug(f"{e=}")
361            if connection:
362                await connection.close()
363            if heartbeat:
364                heartbeat.cancel()
365            self._set_connection(False)
366            self._server_connection = None
367            self._set_status(f"Disconnected: {e.args[0]}")
368            logger.info(f"Connection terminated. {self}")
369            return
370        logger.warning(f"Connection terminated without DisconnectedError {connection}")
371
372    async def _async_connect_local(self):
373        """Wraps `Client._async_connect_remote to cleanup/teardown the local server."""
374        logger.debug("Setting up local server...")
375        assert isinstance(self._server, Server)
376        started = asyncio.Future()
377        server_coro = self._server.async_run(on_start=lambda *a: started.set_result(1))
378        server_task = asyncio.create_task(server_coro)
379        server_task.add_done_callback(lambda *a: self.disconnect())
380        await asyncio.wait((server_task, started), return_when=asyncio.FIRST_COMPLETED)
381        await self._async_connect_remote()
382        self.disconnect()
383        self._server.shutdown()
384
385    async def _handle_handshake(self, connection: ClientConnection):
386        """Handle a new connection's handshake sequence. Modifies the connection object.
387
388        First trade public keys and assign the connection's `tunnel`. Then log in
389        with our username.
390        """
391        # Trade public keys
392        packet = Packet("key_trade", dict(pubkey=self._key.pubkey))
393        response = await connection.send_recv(packet)
394        pubkey = response.payload.get("pubkey")
395        if not pubkey or not isinstance(pubkey, str):
396            raise DisconnectedError("Missing public key string from server.")
397        if self._verify_server_pubkey:
398            if pubkey != self._verify_server_pubkey:
399                raise DisconnectedError("Unverified server public key.")
400            logger.debug(f"Server pubkey verified: {pubkey=}")
401        connection.tunnel = self._key.get_tunnel(pubkey)
402        logger.debug(f"Assigned tunnel: {connection}")
403        # Authenticate
404        handshake_payload = dict(
405            username=self._username,
406            password=self._password,
407            invite_code=self._invite_code,
408        )
409        packet = Packet("handshake", handshake_payload)
410        response = await connection.send_recv(packet)
411        if response.status != Status.OK:
412            m = response.message
413            logger.info(m)
414            raise DisconnectedError(m)
415
416    async def _async_heartbeat(self):
417        """Periodically update while connected and in game.
418
419        Will create a heartbeat request using `Client.heartbeat_payload` and pass the
420        response to `Client.on_heartbeat`.
421        """
422        while True:
423            await asyncio.sleep(self._heartbeat_interval)
424            if self.connected and self.game:
425                packet = Packet(Request.HEARTBEAT_UPDATE, self.heartbeat_payload())
426                self.send(packet, self.on_heartbeat)
427
428    async def _handle_user_connection(self, connection: ClientConnection):
429        """Handle the connection after handshake - process the packet queue."""
430        self._do_disconnect = False
431        while not self._do_disconnect:
432            # Wait for queued packets
433            if not self._packet_queue:
434                await asyncio.sleep(0.01)
435                continue
436            # Send and callback with response
437            packet, callback = self._packet_queue.pop(0)
438            response = await connection.send_recv(packet)
439            if response.status != Status.OK:
440                logger.info(f"Status code: {response.debug_repr}")
441            self._handle_game_change(response)
442            if callback:
443                callback(response)
444            # Disconnect if alerted
445            if response.disconnecting:
446                logger.info(f"Disconnected.\n{packet}\n{response.debug_repr}")
447                raise DisconnectedError(response.message)
448        raise DisconnectedError("Client closed connection.")
449
450    def _set_status(self, status: str, logger: Optional[Callable] = None):
451        """Set the client status message with associated callback."""
452        self._status = status
453        if logger:
454            logger(status)
455        if self.on_status:
456            self.on_status(status)
457
458    def _set_connection(self, set_as: bool, /):
459        """Set the client connection status with associated callback."""
460        if self._connected == set_as:
461            return
462        logger.debug(f"Setting connection as: {set_as}")
463        self._connected = set_as
464        if self.on_connection:
465            self.on_connection(set_as)
466
467    def _handle_game_change(self, response: Response):
468        """Handle game change.
469
470        Set the client game name and heartbeat rate, then call event callback.
471        """
472        game_name = response.game
473        if game_name != self._game:
474            if game_name is None:
475                self._set_status(f"Left game: {self._game}")
476            else:
477                self._set_status(f"Joined game: {game_name}")
478            hb_rate = response.payload.get("heartbeat_rate", DEFAULT_HEARTBEAT_RATE)
479            self._heartbeat_interval = 1 / hb_rate
480            logger.debug(f"{self._heartbeat_interval=}")
481            self._game = game_name
482            if self.on_game:
483                self.on_game(game_name)
484
485    def __repr__(self):
486        """Object repr."""
487        local = "(local) " if self._server else ""
488        conn = "connected" if self.connected else "disconnected"
489        return (
490            f"<{self.__class__.__qualname__} {local}{self._username!r}"
491            f" {conn}, {self._address}:{self._port}"
492            f" @ {id(self):x}>"
493        )

251class Server:
252    """The server that hosts games.
253
254    Subclass from `pgnet.Game` and pass it as the *`game`* argument for the server.
255    Then, use the `Server.async_run` coroutine to start the server.
256
257    By default, the server is configured to listen on localhost. To listen
258    globally, set *`listen_globally`* and *`admin_password`*.
259
260    For games to save and load, *`save_file`* must be set (see also:
261    `pgnet.Game.get_save_string`).
262
263    .. note:: Most home networks require port forwarding to be discoverable by remote
264        clients.
265    """
266
267    def __init__(
268        self,
269        game: Type[Game],
270        /,
271        *,
272        listen_globally: bool = False,
273        port: int = DEFAULT_PORT,
274        admin_password: Optional[str] = None,
275        registration_enabled: bool = True,
276        require_user_password: bool = False,
277        on_connection: Optional[Callable[[str, bool], Any]] = None,
278        verbose_logging: bool = False,
279        save_file: Optional[str | Path] = None,
280    ):
281        """Initialize the server.
282
283        Args:
284            listen_globally: Listen globally instead of localhost only.
285                Requires that *`admin_password`* must be set.
286            port: Port number to listen on.
287            admin_password: Password for admin user with elevated priviliges.
288            registration_enabled: Allow new users to register.
289            require_user_password: Require that users have non-empty passwords.
290            on_connection: Callback for when a username connects or disconnects.
291            verbose_logging: Log *all* packets and responses.
292            save_file: Location of file to save and load server sessions.
293        """
294        if listen_globally and not admin_password:
295            logger.warning(
296                "Created server that listens globally without admin password."
297            )
298        admin_password = admin_password or DEFAULT_ADMIN_PASSWORD
299        self._key: Key = Key()
300        self._stop: Optional[asyncio.Future] = None
301        self._require_user_password = require_user_password
302        self._users: dict[str, User] = dict()
303        self._register_user(ADMIN_USERNAME, admin_password)
304        self._games: dict[str, LobbyGame] = {}
305        self._connections: dict[str, Optional[UserConnection]] = {}
306        self._deleted_users: set[str] = set()
307        self._invite_codes: dict[str, str] = {}
308        self._game_cls: Type[Game] = game
309        self._save_file: Optional[Path] = None if save_file is None else Path(save_file)
310        self._address: str = "" if listen_globally else "localhost"
311        self._port: int = port
312        self.registration_enabled: bool = registration_enabled
313        self.on_connection: Optional[Callable[[str, bool], Any]] = on_connection
314        self.verbose_logging: bool = verbose_logging
315        self._load_from_disk()
316        logger.debug(f"{self._save_file=}")
317        logger.debug(f"{self._game_cls=}")
318        logger.debug(f"{self._key=}")
319        print(f"{admin_password=}")  # Use print instead of logger for password
320
321    async def async_run(self, *, on_start: Optional[Callable] = None) -> int:
322        """Start the server.
323
324        The server will listen for connections and pass them off to the
325        connection handler.
326
327        Args:
328            on_start: Callback for when the server is online and handling messages.
329
330        Returns:
331            Exit code as given by the `shutdown` command. A value of -1 indicates a
332                request to reboot.
333        """
334        if self._stop:
335            raise RuntimeError("Cannot run the server more than once concurrently.")
336        self._stop = asyncio.Future()
337        serving_args = (self._connection_handler, self._address, self._port)
338        try:
339            async with websockets.serve(*serving_args):
340                logger.info(f"Handling messages {self}")
341                if on_start:
342                    on_start()
343                await self._listening_loop(self._stop)
344        except OSError as e:
345            added = OSError(f"Server fail. Perhaps one is already running? {self}")
346            raise added from e
347        self._save_to_disk()
348        result = self._stop.result()
349        logger.info(f"Server stop {result=} {self}")
350        self._stop = None
351        return result
352
353    def shutdown(self, result: int = 0, /):
354        """Stop the server.
355
356        The *result* is passed as the return value (exit code) for `Server.async_run`.
357        """
358        if self._stop and not self._stop.done():
359            self._stop.set_result(result)
360
361    def delete_user(self, username: str):
362        """Disconnect and delete a given username from the server."""
363        if username == ADMIN_USERNAME:
364            logger.warning("Cannot delete admin.")
365            return
366        if username in self._connections:
367            self._deleted_users.add(username)
368        else:
369            self._delete_user(username)
370
371    @property
372    def pubkey(self) -> str:
373        """Public key used for end to end encryption."""
374        return self._key.pubkey
375
376    async def _listening_loop(self, stop_future: asyncio.Future):
377        next_autosave = arrow.now().shift(seconds=AUTOSAVE_INTERVAL)
378        next_interval = arrow.now().shift(seconds=GAME_UPDATE_INTERVAL)
379        while not stop_future.done():
380            await asyncio.sleep(0.1)
381            if arrow.now() >= next_autosave:
382                self._save_to_disk()
383                next_autosave = arrow.now().shift(seconds=AUTOSAVE_INTERVAL)
384            if arrow.now() >= next_interval:
385                for game in self._games.values():
386                    game.update()
387                next_interval = arrow.now().shift(seconds=GAME_UPDATE_INTERVAL)
388
389    async def _connection_handler(self, websocket: ServerWebSocket):
390        """Handle new connections.
391
392        Allows the handshake to fully populate a UserConnection, which may then
393        be handled as a logged in user.
394        """
395        connection = UserConnection(websocket)
396        logger.info(f"New connection: {connection}")
397        try:
398            await self._handle_handshake(connection)
399            self._add_user_connection(connection)
400            logger.info(f"User logged in: {connection}")
401            await self._handle_user(connection)
402        except DisconnectedError as e:
403            logger.debug(f"{e=}")
404        finally:
405            logger.info(f"Closed connection: {connection}")
406            self._remove_user_connection(connection)
407            username = connection.username
408            if username in self._users and username in self._deleted_users:
409                self._delete_user(connection.username)
410
411    async def _handle_handshake(self, connection: UserConnection):
412        """Handle a new connection's handshake sequence. Modifies the connection object.
413
414        First trade public keys and assign the connection's `tunnel`. Then authenticate
415        and assign the connection's `username`.
416        """
417        # Trade public keys
418        packet = await connection.recv()
419        pubkey = packet.payload.get("pubkey")
420        if not pubkey or not isinstance(pubkey, str):
421            response = Response(
422                "Missing public key string.",
423                status=Status.BAD,
424                disconnecting=True,
425            )
426            await connection.send(response)
427            raise DisconnectedError("Incompatible protocol: missing pubkey.")
428        response = Response("key_trade", dict(pubkey=self.pubkey))
429        await connection.send(response)
430        connection.tunnel = self._key.get_tunnel(pubkey)
431        logger.debug(f"Assigned tunnel: {connection}")
432        # Authenticate
433        packet = await connection.recv()
434        username = packet.payload.get("username")
435        password = packet.payload.get("password", "")
436        invite_code = packet.payload.get("invite_code", "")
437        fail = self._check_auth(username, password, invite_code)
438        if fail:
439            # Respond with problem and disconnect
440            response = Response(fail, status=Status.BAD, disconnecting=True)
441            await connection.send(response)
442            raise DisconnectedError("Failed to authenticate.")
443        connection.username = username
444        logger.debug(f"Assigned username: {connection}")
445        if username == ADMIN_USERNAME:
446            logger.warning(f"Authenticated as admin: {connection}")
447        await connection.send(Response("Authenticated."))
448
449    def _check_auth(
450        self,
451        username: str,
452        password: str,
453        invite_code: str,
454    ) -> Optional[str]:
455        """Return failure reason or None."""
456        if not username:
457            return "Missing non-empty username."
458        if username in self._deleted_users:
459            return "User deleted."
460        if username in self._connections:
461            return "Username already connected."
462        if username not in self._users:
463            return self._try_register_user(username, password, invite_code)
464        user = self._users[username]
465        if not user.compare_password(password):
466            return "Incorrect password."
467        return None
468
469    def _try_register_user(
470        self,
471        username: str,
472        password: str,
473        invite_code: str,
474    ) -> Optional[str]:
475        """Return failure reason or None."""
476        if invite_code:
477            invite_username = self._invite_codes.get(invite_code)
478            wrong_code = invite_username is None
479            invite_valid = username == invite_username or invite_username == ""
480            if wrong_code or not invite_valid:
481                return "Incorrect username or invite code."
482        if not (self.registration_enabled or invite_code):
483            return "Registration blocked."
484        if self._require_user_password and not password:
485            return "User password required."
486        if not is_username_allowed(username):
487            return "Username not allowed."
488        self._register_user(username, password)
489        if invite_code:
490            del self._invite_codes[invite_code]
491        return None
492
493    def _register_user(self, username: str, password: str, /):
494        """Register new user."""
495        assert username not in self._users
496        if self._require_user_password and not password:
497            raise ValueError("Server requires password for users.")
498        user = User.from_name_password(username, password)
499        self._users[username] = user
500        logger.info(f"Registered {username=}")
501
502    def _add_user_connection(self, connection: UserConnection):
503        """Add the connection to connected users table."""
504        username = connection.username
505        assert username not in self._connections
506        self._connections[username] = connection
507        if self.on_connection:
508            self.on_connection(username, True)
509
510    def _remove_user_connection(self, connection: UserConnection):
511        """Remove the connection from connected users table if exists."""
512        username = connection.username
513        if username not in self._connections:
514            return
515        self._remove_user_from_game(connection.username)
516        del self._connections[username]
517        if self.on_connection:
518            self.on_connection(username, False)
519
520    async def _handle_user(self, connection: UserConnection):
521        """Handle a logged in user connection - handle packets and return responses."""
522        username = connection.username
523        while True:
524            # Wait for packet from user
525            packet = await connection.recv(timeout=3600.0)
526            # Important! We must set the packet's authenticated username.
527            packet.username = username
528            do_log = self.verbose_logging
529            if do_log:
530                logger.debug(packet)
531            if username in self._deleted_users:
532                response = Response("User deleted.", disconnecting=True)
533            else:
534                response: Response = self._handle_packet(packet)
535            if do_log:
536                logger.debug(f"--> {response}")
537            assert isinstance(response, Response)
538            # Also important, to set the game of the response for the client.
539            response.game = self._connections[username].game
540            await connection.send(response)
541            # The packet handler may have determined we are disconnecting
542            if response.disconnecting:
543                raise DisconnectedError(response.message)
544
545    def _handle_packet(self, packet: Packet) -> Response:
546        """Handle a packet from a logged in user."""
547        # Find builtin handler
548        request_handler = self._request_handlers.get(packet.message)
549        if request_handler:
550            return request_handler(self, packet)
551        # Find game handler
552        game_name: Optional[str] = self._connections[packet.username].game
553        if game_name:
554            return self._handle_game_packet(packet, game_name)
555        # No handler found - not in game and not a builtin request
556        return Response(
557            "Please create/join a game.",
558            self._canned_response_payload | dict(packet=packet.debug_repr),
559            status=Status.UNEXPECTED,
560        )
561
562    def _remove_user_from_game(self, username: str):
563        """Remove user from game and delete the game if expired."""
564        connection = self._connections[username]
565        game = self._games.get(connection.game)
566        if not game:
567            return
568        connection.game = None
569        game.remove_user(username)
570        if game.expired:
571            del self._games[game.name]
572        logger.debug(f"User {username!r} removed from {game}")
573
574    def _delete_user(self, username: str):
575        assert username in self._users and username not in self._connections
576        del self._users[username]
577        if username in self._deleted_users:
578            self._deleted_users.remove(username)
579        logger.info(f"Deleted {username=}")
580
581    @_user_packet_handler()
582    def _handle_game_dir(self, packet: Packet) -> Response:
583        """Create a Response with dictionary of games details."""
584        games_dict = {}
585        for name, game in self._games.items():
586            games_dict[game.name] = dict(
587                name=game.name,
588                users=len(game.connected_users),
589                password_protected=game.password_protected,
590                info=game.get_lobby_info(),
591            )
592        return Response("See payload for games directory.", dict(games=games_dict))
593
594    def _create_game(
595        self,
596        name: str,
597        password: str = "",
598        game_data: Optional[str] = None,
599    ) -> LobbyGame:
600        """Create a new game."""
601        assert name not in self._games
602        game = self._game_cls(name, save_string=game_data)
603        lobbygame = LobbyGame(game, name, password)
604        self._games[name] = lobbygame
605        logger.debug(f"Created game: {lobbygame}")
606        return lobbygame
607
608    def _destroy_game(self, game_name: str):
609        """Destroy an existing game."""
610        game = self._games[game_name]
611        while game.connected_users:
612            self._remove_user_from_game(list(game.connected_users)[0])
613        if game_name in self._games:
614            del self._games[game_name]
615        logger.debug(f"Destroyed game: {game_name!r}")
616
617    @_user_packet_handler()
618    def _handle_join_game(
619        self,
620        packet: Packet,
621        /,
622        *,
623        name: str = "",
624    ) -> Response:
625        """Handle a request to join a game."""
626        game_name = name
627        connection = self._connections[packet.username]
628        current_name: Optional[str] = connection.game
629        if current_name:
630            return Response("Must leave game first.", status=Status.UNEXPECTED)
631        if not game_name:
632            return Response("Please specify a game name.", status=Status.UNEXPECTED)
633        if game_name == current_name:
634            return Response("Already in game.", status=Status.UNEXPECTED)
635        game = self._games.get(game_name)
636        if not game:
637            return self._handle_create_game(packet)
638        password = packet.payload.get("password", "")
639        fail = game.add_user(packet.username, password)
640        if fail:
641            return Response(f"Failed to join game: {fail}", status=Status.UNEXPECTED)
642        connection.game = game_name
643        logger.debug(f"User {packet.username!r} joined: {game}")
644        return Response(
645            f"Joined game: {game_name!r}.",
646            dict(heartbeat_rate=game.heartbeat_rate),
647        )
648
649    @_user_packet_handler()
650    def _handle_leave_game(self, packet: Packet) -> Response:
651        """Handle a request to leave the game."""
652        game_name: Optional[str] = self._connections[packet.username].game
653        if not game_name:
654            return Response("Not in game.", status=Status.UNEXPECTED)
655        self._remove_user_from_game(packet.username)
656        logger.debug(f"User {packet.username!r} left game: {game_name!r}")
657        return Response(f"Left game: {game_name!r}.")
658
659    @_user_packet_handler()
660    def _handle_create_game(
661        self,
662        packet: Packet,
663        /,
664        *,
665        name: str = "",
666        password: str = "",
667    ) -> Response:
668        """Handle request to create a new game specified in the payload."""
669        game_name = name
670        connection = self._connections[packet.username]
671        current_game = connection.game
672        if current_game:
673            return Response("Must leave game first.", status=Status.UNEXPECTED)
674        if not is_gamename_allowed(game_name):
675            return Response("Game name not allowed.", status=Status.UNEXPECTED)
676        if game_name in self._games:
677            return Response("Game name already exists.", status=Status.UNEXPECTED)
678        game = self._create_game(game_name, password)
679        fail = game.add_user(packet.username, password)
680        assert not fail
681        connection.game = game_name
682        logger.debug(f"User {packet.username!r} created game: {game}")
683        return Response(
684            f"Created new game: {game_name!r}.",
685            dict(heartbeat_rate=game.heartbeat_rate),
686        )
687
688    def _handle_game_packet(self, packet: Packet, game_name: str) -> Response:
689        """Routes a packet from a logged in user to the game's packet handler.
690
691        Will use the response's `disconnecting` attribute to remove the user
692        from the game, and then clear the attribute.
693        """
694        game = self._games[game_name]
695        response: Response = game.handle_packet(packet)
696        assert isinstance(response, Response)
697        if response.disconnecting:
698            self._remove_user_from_game(packet.username)
699            response.disconnecting = False
700        return response
701
702    @_user_packet_handler()
703    def _handle_help(self, packet: Packet) -> Response:
704        requests = dict()
705        for name, f in self._request_handlers.items():
706            requests[name] = {
707                name: param.annotation.__name__
708                for name, param in _get_packet_handler_params(f).items()
709            }
710        return Response("See payload for requests.", requests)
711
712    # Admin commands
713    @_user_packet_handler(admin=True)
714    def _admin_shutdown(self, packet: Packet) -> Response:
715        """Shutdown the server."""
716        self.shutdown()
717        return Response("Shutting down...")
718
719    @_user_packet_handler(admin=True)
720    def _admin_create_invite(
721        self,
722        packet: Packet,
723        /,
724        *,
725        username: str = "",
726    ) -> Response:
727        """Create an invite code. Can optionally by for a specific username."""
728        code = os.urandom(2).hex()
729        self._invite_codes[code] = username
730        return Response(f"Created invite code: {code}")
731
732    @_user_packet_handler(admin=True)
733    def _admin_register(
734        self,
735        packet: Packet,
736        /,
737        *,
738        set_as: bool = False,
739    ) -> Response:
740        """Set user registration."""
741        self.registration_enabled = set_as
742        return Response(f"Registration enabled: {set_as}")
743
744    @_user_packet_handler(admin=True)
745    def _admin_delete_user(
746        self,
747        packet: Packet,
748        /,
749        *,
750        username: str = ""
751    ) -> Response:
752        """Delete a user by name."""
753        if username not in self._users:
754            return Response(f"No such username {username!r}")
755        self.delete_user(username)
756        return Response(f"Requested delete user {username!r}")
757
758    @_user_packet_handler(admin=True)
759    def _admin_destroy_game(
760        self,
761        packet: Packet,
762        /,
763        *,
764        name: str = "",
765    ) -> Response:
766        """Destroy a game by name."""
767        game_name = name
768        if game_name not in self._games:
769            return Response(f"No such game: {game_name!r}", status=Status.UNEXPECTED)
770        self._destroy_game(game_name)
771        return Response(f"Destroyed game: {game_name!r}")
772
773    @_user_packet_handler(admin=True)
774    def _admin_save(self, packet: Packet) -> Response:
775        """Save all server data to file."""
776        success = self._save_to_disk()
777        return Response(f"Saved {success=} server data to disk: {self._save_file}")
778
779    @_user_packet_handler(admin=True)
780    def _admin_verbose(
781        self,
782        packet: Packet,
783        /,
784        *,
785        set_as: bool = False,
786    ) -> Response:
787        """Set verbose logging."""
788        self.verbose_logging = set_as
789        return Response(f"Verbose logging enabled: {set_as}")
790
791    @_user_packet_handler(admin=True)
792    def _admin_debug(self, packet: Packet) -> Response:
793        """Return debugging info."""
794        games = [str(game) for name, game in sorted(self._games.items())]
795        connected_users = [str(conn) for u, conn in sorted(self._connections.items())]
796        all_users = sorted(self._users.keys())
797        payload = dict(
798            packet=packet.debug_repr,
799            pubkey=self.pubkey,
800            games=games,
801            connected_users=connected_users,
802            all_users=all_users,
803            registration=self.registration_enabled,
804            invite_codes=self._invite_codes,
805            deleted_users=list(self._deleted_users),
806            verbose=self.verbose_logging,
807        )
808        return Response("Debug", payload)
809
810    @_user_packet_handler(admin=True)
811    def _admin_sleep(self, packet: Packet, /, *, seconds: float = 1) -> Response:
812        """Simulate slow response by blocking for the time specified in payload.
813
814        Warning: this actually blocks the entire server. Time is capped at 5 seconds.
815        """
816        max_sleep = 5
817        seconds = min(max_sleep, seconds)
818        time.sleep(seconds)
819        return Response(f"Slept for {seconds} seconds")
820
821    def _save_to_disk(self) -> bool:
822        """Save all data to disk."""
823        if not self._save_file:
824            return False
825        game_data = []
826        for game in self._games.values():
827            save_string = game.get_save_string()
828            if not save_string:
829                continue
830            game_data.append(dict(
831                name=game.name,
832                password=game.password,
833                data=save_string,
834            ))
835        users = [
836            dict(name=u.name, salt=u.salt, password=u.hashed_password)
837            for u in self._users.values() if u.name != ADMIN_USERNAME
838        ]
839        data = dict(
840            users=users,
841            games=game_data,
842            registration=self.registration_enabled,
843            invite_codes=self._invite_codes,
844        )
845        dumped = json.dumps(data, indent=4)
846        self._save_file.parent.mkdir(parents=True, exist_ok=True)
847        with open(self._save_file, "w") as f:
848            f.write(dumped)
849        logger.debug(
850            f"Saved server data to {self._save_file}"
851            f" ({len(users)} users and {len(game_data)} games)"
852        )
853        return True
854
855    def _load_from_disk(self):
856        if not self._save_file or not self._save_file.is_file():
857            return
858        logger.info(f"Loading server data from {self._save_file}")
859        with open(self._save_file) as f:
860            data = f.read()
861        data = json.loads(data)
862        for user in data["users"]:
863            username = user["name"]
864            if username == ADMIN_USERNAME:
865                continue
866            if not is_username_allowed(username):
867                logger.warning(f"Loaded disallowed {username=}")
868            self._users[username] = u = User(username, user["salt"], user["password"])
869            logger.debug(f"Loaded username: {u!r}")
870        for game in data["games"]:
871            game_name = game["name"]
872            if not is_gamename_allowed(game_name):
873                logger.warning(f"Loaded disallowed {game_name=}")
874            self._create_game(game_name, game["password"], game["data"])
875            logger.debug(f"Loaded game: {self._games[game_name]!r}")
876        self._invite_codes |= data["invite_codes"]
877        self.registration_enabled = data["registration"]
878        logger.debug("Loading disk data complete.")
879
880    def __repr__(self):
881        """Object repr."""
882        address = self._address or "public"
883        return (
884            f"<{self.__class__.__qualname__}"
885            f" serving {address}:{self._port}"
886            f" @ {id(self):x}>"
887        )
888
889    _request_handlers = {
890        Request.HELP: _handle_help,
891        Request.GAME_DIR: _handle_game_dir,
892        Request.CREATE_GAME: _handle_create_game,
893        Request.JOIN_GAME: _handle_join_game,
894        Request.LEAVE_GAME: _handle_leave_game,
895        Request.DEBUG: _admin_debug,
896        Request.SAVE: _admin_save,
897        Request.CREATE_INVITE: _admin_create_invite,
898        Request.DESTROY_GAME: _admin_destroy_game,
899        Request.DELETE_USER: _admin_delete_user,
900        Request.REGISTRATION: _admin_register,
901        Request.VERBOSE: _admin_verbose,
902        Request.SLEEP: _admin_sleep,
903        Request.SHUTDOWN: _admin_shutdown,
904    }
905    _canned_response_payload = dict(commands=list(_request_handlers.keys()))