From ecbcdc9daec52a4746619bdeaf7a75dc8cf11c0a Mon Sep 17 00:00:00 2001
From: Scott Idem <stidem@gmail.com>
Date: Tue, 14 Mar 2023 20:14:58 -0400
Subject: [PATCH] Work on websockets!

---
 app/methods/account_cfg_methods.py |   2 +-
 app/routers/websockets.py          | 226 ++++++++++++++++++++++-------
 2 files changed, 177 insertions(+), 51 deletions(-)

diff --git a/app/methods/account_cfg_methods.py b/app/methods/account_cfg_methods.py
index 32951e9..8efefe5 100644
--- a/app/methods/account_cfg_methods.py
+++ b/app/methods/account_cfg_methods.py
@@ -27,7 +27,7 @@ def load_account_cfg_obj(
     if account_id := redis_lookup_id_random(record_id_random=account_id, table_name='account'): pass
     else: return False
 
-    log.info('Getting Account CFG ID: ${account_id}')
+    log.info(f'Getting Account CFG ID: ${account_id}')
     if account_cfg_rec := sql_select(
         table_name = 'v_account_cfg', # This view should probably be cleaned up
         field_name = 'account_id',
diff --git a/app/routers/websockets.py b/app/routers/websockets.py
index 45344c0..48dc58a 100644
--- a/app/routers/websockets.py
+++ b/app/routers/websockets.py
@@ -2,7 +2,7 @@ from fastapi import APIRouter, FastAPI, Response, WebSocket, WebSocketDisconnect
 from fastapi.responses import HTMLResponse
 from pydantic import BaseModel, EmailStr, Field
 from typing import Dict, List, Optional, Set, Union
-import redis, asyncio, datetime, hashlib, json, os, pathlib, shutil, time
+import redis, asyncio, base64, datetime, hashlib, json, os, pathlib, shutil, time
 
 from app.lib_general import log, logging, common_route_params, Common_Route_Params, common_route_params_min, Common_Route_Params_Min
 from app.config import settings
@@ -29,7 +29,8 @@ html = """
         <script>
             var client_id = Date.now();
             document.querySelector("#ws-id").textContent = client_id;
-            var ws = new WebSocket(`ws://localhost:5005/ws/${client_id}`);
+            var ws = new WebSocket(`ws://fastapi.localhost:8080/ws/${client_id}`);
+            // var ws = new WebSocket(`ws://localhost:5005/ws/${client_id}`);
             // var ws = new WebSocket("ws://localhost:8000/ws");
             // var ws = new WebSocket("ws://fastapi.localhost/ws");
             ws.onmessage = function(event) {
@@ -75,6 +76,45 @@ class ConnectionManager:
     def disconnect(self, websocket: WebSocket):
         self.active_connections.remove(websocket)
 
+    async def echo(self, message: str, websocket: WebSocket):
+        log.setLevel(logging.DEBUG)
+        # log.debug(dir(websocket))
+        log.debug(vars(websocket))
+        log.debug(websocket.url)
+        log.debug(websocket.client)
+        log.debug(websocket.client_state)
+        log.debug(websocket.headers['sec-websocket-key'])
+        # log.debug(base64.decode(bytes(websocket.headers['sec-websocket-key']), 'utf-8'))
+        await websocket.send_text(message)
+
+    async def direct(self, from_client_id: str, to_client_id: str, data: dict):
+        log.setLevel(logging.DEBUG)
+        for connection in self.active_connections:
+            log.debug(vars(connection))
+            log.debug(connection)
+            await connection.send_text(message)
+
+    async def group(self, group_id: str, data: str):
+        log.setLevel(logging.DEBUG)
+        log.debug(locals())
+
+        for connection in self.active_connections:
+            log.debug(vars(connection))
+            # websocket.path_params.get('client_id')
+            # if connection.scope.get('path') == group_id:
+            if connection.path_params.get('group_id') == group_id:
+                log.info('Found matching Group ID')
+                await connection.send_json(data)
+
+    # NOTE: Same as group, but no filter based on path
+    async def broadcast(self, message: str):
+        log.setLevel(logging.INFO)
+        log.debug(locals())
+
+        for connection in self.active_connections:
+            log.debug(vars(connection))
+            await connection.send_text(message)
+
     async def send_personal_message(self, message: str, websocket: WebSocket):
         log.setLevel(logging.DEBUG)
         # log.debug(dir(websocket))
@@ -83,13 +123,9 @@ class ConnectionManager:
         log.debug(websocket.client)
         log.debug(websocket.client_state)
         log.debug(websocket.headers['sec-websocket-key'])
+        # log.debug(base64.decode(bytes(websocket.headers['sec-websocket-key']), 'utf-8'))
         await websocket.send_text(message)
 
-    async def broadcast(self, message: str):
-        log.setLevel(logging.DEBUG)
-        for connection in self.active_connections:
-            log.debug(connection)
-            await connection.send_text(message)
 
 
 manager = ConnectionManager()
@@ -101,61 +137,100 @@ manager = ConnectionManager()
 # /person/<id> (for one specific person; handles send and receiving their messages)
 
 
-@router.websocket('/ws/room/{room_id}')
-async def ws_room_id(
+@router.websocket('/ws/client/{client_id}')
+async def ws_client_id(
         websocket: WebSocket,
-        room_id: str,
+        client_id: str,
         ):
     await manager.connect(websocket)
-    await manager.broadcast(f'Welcome to room "{room_id}"!')
     try:
         while True:
-            data = await websocket.receive_json()
-            log.debug(data)
-            data_dict = data
-            # data_dict = json.loads(data)
-            log.debug(data_dict['client_id'])
-            client_id = data_dict['client_id']
-            await manager.send_personal_message(f'You wrote: {data}', websocket)
-            await manager.broadcast(f'Client #{client_id} says: {data}')
-    except WebSocketDisconnect:
-        manager.disconnect(websocket)
-        await manager.broadcast(f'Client left the room')
-
-
-    # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
-    # time.sleep(3.5) # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
-    # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
-
-
-
-
-@router.websocket('/ws/looping')
-async def ws_looping(
-        websocket: WebSocket,
-        ):
-    await manager.connect(websocket)
-    # await manager.broadcast(f'Welcome to looping')
-    try:
-        while True:
-            # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
-            # await time.sleep(3.5) # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
-            # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
-            # data = await websocket.receive_json()
+            data = await websocket.receive_json() # Returns dict
             # log.debug(data)
-            # data_dict = data
-            # data_dict = json.loads(data)
-            # log.debug(data_dict['client_id'])
-            # await manager.send_personal_message(f'You wrote: {data}', websocket)
-            await manager.broadcast(f'Loop!!!')
+
+            # echo (echo message), dm (direct message), group (group message), all (broadcast message to all), cmd, group_cmd(?)
+            msg_type = data.get('type')
+            cmd = data.get('cmd')
+            msg = data.get('msg')
+            to_client_id = data.get('to_client_id')
+            to_group_id = data.get('to_group_id')
+
+            log.setLevel(logging.INFO)
+            log.info(f'Client ID: {client_id}; Type: {msg_type};')
+            log.debug(f'Command: {cmd}')
+            log.debug(f'Message: {msg}')
+            log.debug(f'To Client ID: {to_client_id}')
+            log.debug(f'To Group ID: {to_group_id}')
+
+            if msg_type:
+                if msg_type == 'echo':
+                    await manager.echo(f'Echo: {data}', websocket)
+                elif msg_type == 'dm':
+                    await manager.direct(from_client_id=client_id, to_client_id=to_client_id, data=data)
+                elif msg_type == 'group':
+                    await manager.broadcast(f'Group: {data}')
+                elif msg_type == 'all':
+                    await manager.broadcast(f'All: {data}')
+                else:
+                    await manager.broadcast(f'Unknown: {data}')
+            else:
+                await manager.broadcast(f'MSG: {data}')
+
     except WebSocketDisconnect:
         manager.disconnect(websocket)
-        await manager.broadcast(f'Client left looping')
+        # await manager.broadcast(f'Client #{client_id} left')
+
+
+@router.websocket('/ws/group/{group_id}/client/{client_id}')
+async def ws_client_id(
+        websocket: WebSocket,
+        group_id: str,
+        client_id: str,
+        ):
+    await manager.connect(websocket)
+    try:
+        while True:
+            data = await websocket.receive_json() # Returns dict
+            # log.debug(data)
+
+            # group_path_id = f'/ws/group/{group_id}'
+            # client_id = data.get('client_id')
+            # echo (echo message), dm (direct message), group (group message), all (broadcast message to all), cmd, group_cmd(?)
+            msg_type = data.get('type')
+            cmd = data.get('cmd')
+            msg = data.get('msg')
+
+            log.setLevel(logging.INFO)
+            log.info(f'Group ID: {group_id}; Client ID: {client_id}; Type: {msg_type};')
+            log.debug(f'Command: {cmd}')
+            log.debug(f'Message: {msg}')
+
+            await manager.group(group_id=group_id, data=data)
+
+            # if msg_type:
+            #     if msg_type == 'echo':
+            #         await manager.echo(f'Echo: {data}', websocket)
+            #     elif msg_type == 'dm':
+            #         await manager.direct(f'DM: {msg}')
+            #     elif msg_type == 'group':
+            #         await manager.group(group_id=group_id, f'Group: {data}')
+            #     elif msg_type == 'all':
+            #         await manager.broadcast(f'All: {data}')
+            #     elif msg_type == 'cmd':
+            #         await manager.broadcast(f'Command: {data}')
+            #     else:
+            #         await manager.broadcast(f'Unknown: {data}')
+            # else:
+            #     await manager.broadcast(f'MSG: {data}')
+
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        # await manager.broadcast(f'Client #{client_id} left')
 
 
 
 @router.websocket('/ws/{client_id}')
-async def ws_client_id(
+async def ws_id(
         websocket: WebSocket,
         client_id: int,
         ):
@@ -174,6 +249,57 @@ async def ws_client_id(
         await manager.broadcast(f'Client #{client_id} left')
 
 
+
+# @router.websocket('/ws/room/{room_id}')
+# async def ws_room_id(
+#         websocket: WebSocket,
+#         room_id: str,
+#         ):
+#     await manager.connect(websocket)
+#     await manager.broadcast(f'Welcome to room "{room_id}"!')
+#     try:
+#         while True:
+#             data = await websocket.receive_json()
+#             log.debug(data)
+#             data_dict = data
+#             # data_dict = json.loads(data)
+#             log.debug(data_dict['client_id'])
+#             client_id = data_dict['client_id']
+#             await manager.send_personal_message(f'You wrote: {data}', websocket)
+#             await manager.broadcast(f'Client #{client_id} says: {data}')
+#     except WebSocketDisconnect:
+#         manager.disconnect(websocket)
+#         await manager.broadcast(f'Client left the room')
+
+
+#     # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+#     # time.sleep(3.5) # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+#     # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+
+
+# @router.websocket('/ws/looping')
+# async def ws_looping(
+#         websocket: WebSocket,
+#         ):
+#     await manager.connect(websocket)
+#     # await manager.broadcast(f'Welcome to looping')
+#     try:
+#         while True:
+#             # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+#             # await time.sleep(3.5) # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+#             # NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING NOTE: WARNING
+#             # data = await websocket.receive_json()
+#             # log.debug(data)
+#             # data_dict = data
+#             # data_dict = json.loads(data)
+#             # log.debug(data_dict['client_id'])
+#             # await manager.send_personal_message(f'You wrote: {data}', websocket)
+#             await manager.broadcast(f'Loop!!!')
+#     except WebSocketDisconnect:
+#         manager.disconnect(websocket)
+#         await manager.broadcast(f'Client left looping')
+
+
 # @router.websocket("/ws/{client_id}")
 # async def websocket_endpoint(
 #         websocket: WebSocket,