Add websocket notif/ops and process appropriately.

Author: EvieePy

api/server.py

@@ -35,6 +35,7 @@
 from .middleware.auth import AuthBackend
 from .routes.applications import Applications
 from .routes.auth import Auth
+from .routes.members import Members
 from .routes.users import Users
 
 
@@ -43,15 +44,15 @@ def __init__(self, *, session: aiohttp.ClientSession, database: core.Database) -
         self.session = session
         self.database = database
 
-        views: list[core.View] = [Users(self), Auth(self), Applications(self)]
+        views: list[core.View] = [Users(self), Auth(self), Applications(self), Members(self)]
         middleware: list[Middleware] = [
             Middleware(CORSMiddleware, allow_origins=['*'], allow_methods=['*'], allow_headers=['*']),
             Middleware(AuthenticationMiddleware, backend=AuthBackend(self)),
         ]
 
         self.sockets: dict[int, WebSocket] = {}
         self.subscription_sockets: dict[str, set[int]] = {
-            'discord_py_mod_log': set()
+            core.WebsocketSubscriptions.DPY_MOD_LOG: set()
         }
 
         super().__init__(
@@ -82,7 +83,7 @@ async def websocket_connector(self, websocket: WebSocket) -> None:
 
         # Send the initial accepted response. Includes user_id and subscriptions... op: 0
         data: dict[str, Any] = {
-            'op': core.WebsocketOPCodes.ACCEPTED,
+            'op': core.WebsocketOPCodes.HELLO,
             'user_id': uid,
             'subscriptions': subscriptions
         }
@@ -93,14 +94,73 @@ async def websocket_connector(self, websocket: WebSocket) -> None:
         while True:
 
             try:
-                message: dict[str | int, Any] = await websocket.receive_json()
+                message: dict[str, Any] = await websocket.receive_json()
             except WebSocketDisconnect:
                 break
 
-            # TODO: Process messages...
-            print(message)
+            op: str | None = message.get('op')
+
+            if op == core.WebsocketOPCodes.SUBSCRIBE:
+                response = self.websocket_subscribe(uid=uid, message=message)
+                await websocket.send_json(data=response)
+
+            elif op == core.WebsocketOPCodes.UNSUBSCRIBE:
+                response = self.websocket_unsubscribe(uid=uid, message=message)
+                await websocket.send_json(data=response)
+
+            else:
+                response = {
+                    'op': core.WebsocketOPCodes.NOTIFICATION,
+                    'type': core.WebsocketNotificationTypes.UNKNOWN_OP,
+                    'received': op
+                }
+                await websocket.send_json(data=response)
 
         # Remove the websocket and it's subscriptions...
         del self.sockets[uid]
         for sub in subscriptions:
             self.subscription_sockets[sub].remove(uid)
+
+    def websocket_subscribe(self, *, uid: int, message: dict[str, Any]) -> dict[str, Any]:
+        subs: list[str] = message.get('subscriptions', [])
+
+        # Filter out bad subscriptions...
+        valid: list[str] = list(self.subscription_sockets.keys())
+        subscriptions: list[str] = [sub for sub in subs if sub in valid]
+
+        for sub in subscriptions:
+            self.subscription_sockets[sub].add(uid)
+
+        subscribed: list[str] = [sub for sub in self.subscription_sockets if uid in self.subscription_sockets[sub]]
+
+        data: dict[str, Any] = {
+            'op': core.WebsocketOPCodes.NOTIFICATION,
+            'type': core.WebsocketNotificationTypes.SUBSCRIPTION_ADDED,
+            'user_id': uid,
+            'added': subscriptions,
+            'subscriptions': subscribed
+        }
+
+        return data
+
+    def websocket_unsubscribe(self, *, uid: int, message: dict[str, Any]) -> dict[str, Any]:
+        subs: list[str] = message.get('subscriptions', [])
+
+        # Filter out bad subscriptions...
+        valid: list[str] = list(self.subscription_sockets.keys())
+        subscriptions: list[str] = [sub for sub in subs if sub in valid]
+
+        for sub in subscriptions:
+            self.subscription_sockets[sub].remove(uid)
+
+        subscribed: list[str] = [sub for sub in self.subscription_sockets if uid in self.subscription_sockets[sub]]
+
+        data: dict[str, Any] = {
+            'op': core.WebsocketOPCodes.NOTIFICATION,
+            'type': core.WebsocketNotificationTypes.SUBSCRIPTION_REMOVED,
+            'user_id': uid,
+            'removed': subscriptions,
+            'subscriptions': subscribed
+        }
+
+        return data

core/utils.py

@@ -32,7 +32,15 @@
 from starlette.types import Receive, Scope, Send
 
 
-__all__ = ('route', 'View', 'Application', 'WebsocketCloseCodes', 'WebsocketOPCodes')
+__all__ = (
+    'route',
+    'View',
+    'Application',
+    'WebsocketCloseCodes',
+    'WebsocketOPCodes',
+    'WebsocketSubscriptions',
+    'WebsocketNotificationTypes'
+)
 
 ResponseType: TypeAlias = Coroutine[Any, Any, Response]
 
@@ -213,9 +221,30 @@ class WebsocketCloseCodes:
 
     NORMAL: int = 1000
     ABNORMAL: int = 1006
-    INVALID_SUBSCRIPTION: int = 4005
 
 
 class WebsocketOPCodes:
 
-    ACCEPTED: int = 0
+    # Sent...
+    HELLO: int = 0
+    EVENT: int = 1
+    NOTIFICATION: int = 2
+
+    # Received...
+    SUBSCRIBE: str = 'subscribe'
+    UNSUBSCRIBE: str = 'unsubscribe'
+
+
+class WebsocketSubscriptions:
+
+    DPY_MOD_LOG: str = 'discord_py_mod_log'
+
+
+class WebsocketNotificationTypes:
+
+    # Subscriptions...
+    SUBSCRIPTION_ADDED: str = 'subscription_added'
+    SUBSCRIPTION_REMOVED: str = 'subscription_removed'
+
+    # Failures...
+    UNKNOWN_OP: str = 'unknown_op'