pgnet.util

  1"""Common utilities for pgnet."""
  2
  3from typing import Optional, Callable
  4from loguru import logger
  5import asyncio
  6import enum
  7import functools
  8import arrow
  9import json
 10from dataclasses import dataclass, field
 11import websockets
 12from websockets.legacy.protocol import WebSocketCommonProtocol as WebSocket
 13import nacl.public
 14from nacl.encoding import Base64Encoder
 15
 16
 17DEFAULT_PORT = 38929
 18ADMIN_USERNAME = "admin"
 19DEFAULT_ADMIN_PASSWORD = "admin"
 20
 21
 22class Request:
 23    """Strings used as a message in `Packet` for common requests.
 24
 25    It is recommended to use the `pgnet.Client` API instead of these. `Game` classes
 26    should avoid using the string values here as messages in a `Packet` - such packets
 27    will be specially handled by the server.
 28
 29    In Python 3.11, this should be an `enum.StrEnum`.
 30    """
 31
 32    # Normal requests
 33    HELP: str = "__pgnet__.help"
 34    """Request available commands."""
 35    GAME_DIR: str = "__pgnet__.game_dir"
 36    """Request the games directory."""
 37    JOIN_GAME: str = "__pgnet__.join_game"
 38    """Request to join a game."""
 39    LEAVE_GAME: str = "__pgnet__.leave_game"
 40    """Request to leave a game."""
 41    CREATE_GAME: str = "__pgnet__.create_game"
 42    """Request to create and join a game."""
 43    HEARTBEAT_UPDATE: str = "__pgnet__.heartbeat_update"
 44    """Request a heartbeat update from the game."""
 45
 46    # Admin requests
 47    SHUTDOWN: str = "__pgnet__.shutdown"
 48    """Shutdown the server."""
 49    REGISTRATION: str = "__pgnet__.set_registration"
 50    """Configure registration."""
 51    CREATE_INVITE: str = "__pgnet__.create_invite"
 52    """Create an invite code."""
 53    DELETE_USER: str = "__pgnet__.delete_user"
 54    """Kick a username."""
 55    DESTROY_GAME: str = "__pgnet__.destroy_game"
 56    """Destroy a game."""
 57    SAVE: str = "__pgnet__.save"
 58    """Save server data to disk."""
 59    VERBOSE: str = "__pgnet__.set_verbose"
 60    """Configure logging verbosity."""
 61    DEBUG: str = "__pgnet__.debug"
 62    """Get debug info."""
 63    SLEEP: str = "__pgnet__.sleep"
 64    """Block entire server."""
 65
 66
 67@enum.unique
 68class Status(enum.IntEnum):
 69    """Integer status codes for a `Response` to client requests.
 70
 71    These are used internally by `pgnet.Server`. It is possible but not required to use
 72    these in games and clients.
 73    """
 74
 75    OK: int = 0
 76    """Indicates success without issues."""
 77    BAD: int = 1
 78    """Indicates fatal error."""
 79    UNEXPECTED: int = 2
 80    """Indicates an issue."""
 81
 82
 83class DisconnectedError(Exception):
 84    """Raised when a connection has been or should be closed.
 85
 86    The cause of the error is passed as the first argument as a string.
 87    """
 88    pass
 89
 90
 91class CipherError(Exception):
 92    """Raised when failing to encrypt plaintext or decrypt ciphertext."""
 93    pass
 94
 95
 96@dataclass
 97class Packet:
 98    """The dataclass used to send messages from client to server.
 99
100    Clients need only concern themselves with `Packet.message` and `Packet.payload`.
101    """
102
103    message: str
104    """Message text."""
105    payload: dict = field(default_factory=dict, repr=False)
106    """Dictionary of arbitrary data. Must be JSON-able."""
107    created_on: Optional[str] = None
108    """The creation time of the packet."""
109    username: Optional[str] = field(default=None, repr=False)
110    """Used by the server for identification.
111
112    Setting this on the client side has no effect.
113    """
114
115    def __post_init__(self):
116        """Set creation date."""
117        self.created_on = self.created_on or arrow.now().for_json()
118
119    def serialize(self) -> str:
120        """Convert into a string."""
121        data = {k: getattr(self, k) for k in self.__dataclass_fields__.keys()}
122        try:
123            return json.dumps(data)
124        except Exception as e:
125            m = f"Failed to serialize. See above exception.\n{self.debug_repr}"
126            raise TypeError(m) from e
127
128    @classmethod
129    def deserialize(cls, raw_data: str, /) -> "Packet":
130        """Convert a string into a Packet."""
131        try:
132            data = json.loads(raw_data)
133        except Exception as e:
134            m = f"Failed to deserialize. See above exception.\n{raw_data=}"
135            raise TypeError(m) from e
136        return cls(**data)
137
138    @property
139    def debug_repr(self) -> str:
140        """Repr with more data."""
141        return f"{self!r}+{self.payload}"
142
143
144@dataclass
145class Response:
146    """The dataclass used to respond from server to client.
147
148    Games and Clients need only concern themselves with `Response.message`,
149    `Response.payload`, and `Response.status`.
150    """
151
152    message: str
153    """Message text."""
154    payload: dict = field(default_factory=dict, repr=False)
155    """Dictionary of arbitrary data. Must be JSON-able."""
156    status: int = Status.OK
157    """`Status` code for handling the request that this is responding to."""
158    created_on: Optional[str] = None
159    """The creation time of the packet."""
160    disconnecting: bool = field(default=False, repr=False)
161    """Used by the server to notify the client that the connection is being closed."""
162    game: Optional[str] = field(default=None, repr=False)
163    """Used by the server to notify the client of their current game name."""
164
165    def __post_init__(self):
166        """Set creation date."""
167        self.created_on = self.created_on or arrow.now().for_json()
168
169    def serialize(self) -> str:
170        """Convert into a string."""
171        data = {k: getattr(self, k) for k in self.__dataclass_fields__.keys()}
172        try:
173            return json.dumps(data)
174        except Exception as e:
175            m = f"Failed to serialize. See above exception.\n{self.debug_repr}"
176            raise TypeError(m) from e
177
178    @classmethod
179    def deserialize(cls, raw_data: str, /) -> "Response":
180        """Convert a string into a Response."""
181        try:
182            data = json.loads(raw_data)
183        except Exception as e:
184            m = f"Failed to deserialize. See above exception.\n{raw_data=}"
185            raise TypeError(m) from e
186        return cls(**data)
187
188    @property
189    def debug_repr(self) -> str:
190        """Repr with more data."""
191        return (
192            f"{self!r}(game={self.game!r}, disconnecting={self.disconnecting})"
193            f"+{self.payload}"
194        )
195
196
197class Tunnel:
198    """A shared key for end to end encryption of strings.
199
200    Given a personal private key and an external public key, will encrypt
201    and decrypt string messages. Tunnels should be obtained via a `Key` class.
202    See `Key` class  documentation for more details.
203    """
204    def __init__(self, priv: nacl.public.PrivateKey, pubkey: str):
205        """See class documentation for more details."""
206        pub = nacl.public.PublicKey(pubkey.encode(), encoder=Base64Encoder)
207        self._box = nacl.public.Box(priv, pub)
208        self._pubkey = pubkey  # For the repr
209
210    def encrypt(self, plaintext: str) -> str:
211        """Encrypt a message."""
212        try:
213            encrypted_message = self._box.encrypt(
214                plaintext.encode(),
215                encoder=Base64Encoder,
216            )
217            ciphertext = encrypted_message.decode()
218            return ciphertext
219        except Exception as e:
220            raise CipherError("Failed to encrypt the plaintext.") from e
221
222    def decrypt(self, ciphertext: str) -> str:
223        """Decrypt a message."""
224        try:
225            ciphertext = ciphertext.encode()
226            plaintext = self._box.decrypt(ciphertext, encoder=Base64Encoder).decode()
227            return plaintext
228        except Exception as e:
229            raise CipherError("Failed to decrypt the ciphertext.") from e
230
231    def __repr__(self) -> str:
232        """Object repr."""
233        return f"<{self.__class__.__qualname__} pubkey={self._pubkey[:6]}...>"
234
235
236class Key:
237    """A key manager for end to end encryption of strings.
238
239    Each instance of this class represents a private-public key pair. It can also
240    generate a `Tunnel` (shared key) with other public keys by using `Key.get_tunnel`.
241    These let you encrypt and decrypt messages with other public keys.
242
243    The `Key` and `Tunnel` classes are thin wrappers around `pynacl`'s `PrivateKey` and
244    `Box` classes that can handle public keys and messages as strings.
245    """
246
247    def __init__(self):
248        """See class documentation for details."""
249        self._private_key = nacl.public.PrivateKey.generate()
250        self._public_key = self._private_key.public_key
251        self._tunnels: dict[str, Tunnel] = {}
252
253    @property
254    def pubkey(self) -> str:
255        """The public key in string format."""
256        return self._public_key.encode(encoder=Base64Encoder).decode()
257
258    def get_tunnel(self, pubkey: str) -> Tunnel:
259        """Get a `Tunnel` for a given pubkey."""
260        tunnel = self._tunnels.get(pubkey)
261        if not tunnel:
262            tunnel = self._tunnels[pubkey] = Tunnel(self._private_key, pubkey)
263        return tunnel
264
265    def __repr__(self) -> str:
266        """Object repr."""
267        return f"<{self.__class__.__qualname__} pubkey={self.pubkey}>"
268
269
270@dataclass
271class Connection:
272    """Wrapper for `websocket` with end to end encryption.
273
274    Methods will raise a `DisconnectedError` with the reason for failing to communicate
275    with the websocket.
276    """
277
278    websocket: WebSocket = field(repr=False)
279    """The websocket of this connection."""
280    tunnel: Optional[Tunnel] = None
281    """The `Tunnel` assigned for this connection.
282
283    If set, packets sent and responses received will be encrypted and decrypted using
284    the tunnel.
285    """
286    remote: str = ""
287    """Full address of the other end of the connection."""
288
289    def __post_init__(self):
290        """Set the remote from the websocket."""
291        address, port, *_ = self.websocket.remote_address
292        if address == "::1":
293            address = "localhost"
294        self.remote = f"{address}:{port}"
295
296    @staticmethod
297    def _exception_disconnect(f: Callable, /):
298        """Decorator raising `DisconnecetedError` on connection or cipher exceptions."""
299        @functools.wraps(f)
300        async def wrapper(*args, **kwargs):
301            try:
302                return await f(*args, **kwargs)
303            except websockets.exceptions.ConnectionClosed as e:
304                raise DisconnectedError("Connection closed.") from e
305            except asyncio.exceptions.TimeoutError as e:
306                raise DisconnectedError("Connection timed out.") from e
307            except CipherError as e:
308                raise DisconnectedError("Failed to encrypt message.") from e
309        return wrapper
310
311    @_exception_disconnect
312    async def send(self, message: str, /, *, timeout: float = 5.0):
313        """Send a string to `Connection.websocket`.
314
315        Will be encrypted using `Connection.tunnel` if set.
316        """
317        if self.tunnel:
318            message = self.tunnel.encrypt(message)
319        await asyncio.wait_for(self.websocket.send(message), timeout=timeout)
320
321    @_exception_disconnect
322    async def recv(self, *, timeout: float = 5.0) -> str:
323        """Receive a string from `Connection.websocket`.
324
325        Will be decrypted using `Connection.tunnel` if set.
326        """
327        message: str = await asyncio.wait_for(self.websocket.recv(), timeout=timeout)
328        if self.tunnel:
329            message = self.tunnel.decrypt(message)
330        return message
331
332    async def close(self):
333        """Close the websocket."""
334        await self.websocket.close()
335
336
337class Game:
338    """Subclass to implement game logic.
339
340    This class should not be initialized directly, it is initialized by the server.
341    """
342
343    persistent: bool = False
344    """Set as persistent to allow the game to persist even without players.
345
346    This is required for saving and loading (see also: `Game.get_save_string`).
347    """
348    heartbeat_rate: float = 10
349    """How many times per second the client should check for updates.
350
351    See `Game.handle_heartbeat`.
352    """
353
354    def __init__(self, name: str, save_string: Optional[str] = None):
355        """The server initializes this class for every game started.
356
357        Args:
358            name: Game instance name.
359            save_string: Game data loaded from disk from last server session as given by
360                `Game.get_save_string`.
361        """
362        pass
363
364    def user_joined(self, username: str):
365        """Called when a user joins the game."""
366        pass
367
368    def user_left(self, username: str):
369        """Called when a user leaves the game."""
370        pass
371
372    def handle_packet(self, packet: Packet) -> Response:
373        """Packet handling for heartbeat updates and game requests.
374
375        Most use cases should override `Game.handle_game_packet` and
376        `Game.handle_heartbeat` instead of this method.
377        """
378        if packet.message == Request.HEARTBEAT_UPDATE:
379            return self.handle_heartbeat(packet)
380        return self.handle_game_packet(packet)
381
382    def handle_game_packet(self, packet: Packet) -> Response:
383        """Override this method to implement packet handling.
384
385        See also: `pgnet.Client.send`.
386        """
387        return Response(
388            f"No packet handling configured for {self.__class__.__qualname__}",
389            status=Status.UNEXPECTED,
390        )
391
392    def handle_heartbeat(self, packet: Packet) -> Response:
393        """Override this method to implement heartbeat updates.
394
395        See also: `Game.heartbeat_rate`, `pgnet.Client.heartbeat_payload`,
396        `pgnet.Client.on_heartbeat`.
397        """
398        return Response(
399            f"No heartbeat update configured for {self.__class__.__qualname__}"
400        )
401
402    def get_save_string(self) -> str:
403        """Override this method to save game data to disk.
404
405        If `Game.persistent`, this method is called by the server periodically and when
406        shutting down. In the next session, the server will recreate the game with this
407        string passed as *`save_string`* to `Game.__init__`.
408
409        .. note:: `pgnet.Server` must be configured to enable saving and loading.
410        """
411        return ""
412
413    def get_lobby_info(self) -> str:
414        """Override this method to post public game info in lobby.
415
416        This will be used to give users more information about a game before joining.
417        """
418        return ""
419
420    def update(self):
421        """Called on an interval by the server.
422
423        Override this method to implement background game logic tasks.
424        """
425        pass
426
427
428def enable_logging(enable: bool = True, /):
429    """Enable/disable logging from the PGNet library."""
430    if enable:
431        logger.enable("pgnet")
432    else:
433        logger.disable("pgnet")
434
435
436__all__ = (
437    "Game",
438    "Packet",
439    "Response",
440    "Status",
441    "Request",
442    "Tunnel",
443    "Key",
444    "Connection",
445    "CipherError",
446    "DisconnectedError",
447    "enable_logging",
448)