chat works again

nightwatch/server/__init__.py

@@ -13,23 +13,22 @@
 # Handle state
 class NightwatchStateManager():
     def __init__(self) -> None:
-        self.clients = set()
+        self.clients = {}
         self.message_buffer = []
 
     def add_client(self, client: WebSocketCommonProtocol) -> None:
-        self.clients.add(client)
+        self.clients[client] = None
 
     def remove_client(self, client: WebSocketCommonProtocol) -> None:
-        self.clients.remove(client)
+        if client in self.clients:
+            del self.clients[client]
 
 state = NightwatchStateManager()
 
 # Socket entrypoint
 async def connection(websocket: WebSocketCommonProtocol) -> None:
     try:
-        state.add_client(websocket)
-
-        client = NightwatchClient(websocket)
+        client = NightwatchClient(state, websocket)
         async for message in websocket:
             message = orjson.loads(message)
             if message.get("type") not in registry.commands:

nightwatch/server/utils/commands.py

@@ -3,9 +3,19 @@
 # Modules
 from typing import Callable
 
+import orjson
+import websockets
+
 from . import models
 from .websocket import NightwatchClient
 
+from nightwatch.config import config
+
+# Constants
+class Constant:
+    SERVER_USER: dict[str, str] = {"name": "Nightwatch", "color": "gray"}
+    SERVER_NAME: str = config["server.name"] or "Untitled Server"
+
 # Handle command registration
 class CommandRegistry():
     def __init__(self) -> None:
@@ -22,4 +32,30 @@ def callback(function: Callable) -> None:
 # Setup commands
 @registry.command("identify")
 async def command_identify(state, client: NightwatchClient, data: models.IdentifyModel) -> None:
-    print(data)
+    if client.identified:
+        return await client.send("error", text = "You have already identified.")
+
+    elif data.name.lower() in ["nightwatch", "admin", "administrator", "moderator"]:
+        return await client.send("error", text = "The specified username is reserved.")
+
+    elif data.name in state.clients.values():
+        return await client.send("error", text = "Specified username is already taken.")
+
+    client.set_user_data(data.model_dump())
+    client.identified = True
+
+    await client.send("server", name = Constant.SERVER_NAME, online = len(state.clients))
+    websockets.broadcast(state.clients.keys(), orjson.dumps({
+        "type": "message",
+        "data": {"text": f"{data.name} joined the chatroom.", "user": Constant.SERVER_USER}
+    }))
+
+@registry.command("message")
+async def command_message(state, client: NightwatchClient, data: models.MessageModel) -> None:
+    if not client.identified:
+        return await client.send("error", text = "You must identify before sending a message.")
+
+    websockets.broadcast(state.clients.keys(), orjson.dumps({
+        "type": "message",
+        "data": {"text": data.text, "user": client.user_data}
+    }))

nightwatch/server/utils/models.py

@@ -2,10 +2,13 @@
 
 # Modules
 from typing import Annotated
-from pydantic import BaseModel, StringConstraints
+from pydantic import BaseModel, PlainSerializer, StringConstraints
 from pydantic_extra_types.color import Color
 
 # Models
 class IdentifyModel(BaseModel):
     name: Annotated[str, StringConstraints(min_length = 3, max_length = 32)]
-    color: Color
+    color: Annotated[Color, PlainSerializer(lambda c: c.as_hex(), return_type = str, when_used = "always")]
+
+class MessageModel(BaseModel):
+    text: Annotated[str, StringConstraints(min_length = 1, max_length = 300)]

nightwatch/server/utils/websocket.py

@@ -1,14 +1,25 @@
 # Copyright (c) 2024 iiPython
 
 # Modules
+from typing import Any
+
 import orjson
 from websockets import WebSocketCommonProtocol
 
 class NightwatchClient():
     """This class acts as a wrapper on top of WebSocketCommonProtocol that implements
     data serialization through orjson."""
-    def __init__(self, client: WebSocketCommonProtocol) -> None:
+    def __init__(self, state, client: WebSocketCommonProtocol) -> None:
         self.client = client
+        self.identified = False
+
+        self.state = state
+        self.state.add_client(client)
 
     async def send(self, message_type: str, **message_data) -> None:
         await self.client.send(orjson.dumps({"type": message_type, "data": message_data}))
+
+    # Handle user data (ie. name and color)
+    def set_user_data(self, data: dict[str, Any]) -> None:
+        self.user_data = data
+        self.state.clients[self.client] = data["name"]