Rename files

transphase/transphase.py

@@ -1,58 +1,33 @@
-#!/usr/bin/env python
-import asyncio
-import websockets
-import logging 
 import json
-
-
-logging.basicConfig(
-    format="%(asctime)s %(message)s",
-    level=logging.INFO,
-)
+import asyncio
+from quart import Quart, websocket
 
 peer_list = {}
 
-class LoggerAdapter(logging.LoggerAdapter):
-    """Add connection ID and client IP address to websockets logs."""
-    def process(self, msg, kwargs):
-        try:
-            websocket = kwargs["extra"]["websocket"]
-        except KeyError:
-            return msg, kwargs
-        rip = websocket.remote_address[0]
-        try: 
-            xff = websocket.request_headers.get("X-Forwarded-For")
-        except:
-            xff = "None"
-        return f"{websocket.id} {rip} {xff}", kwargs
+app = Quart(__name__)
 
-async def handler(websocket):
+@app.route("/")
+async def retmain():
+    return f"Ready to relay\n"
+
+@app.websocket("/ws")
+async def handle_websockets():
     peer_group_id = None
-    try:
-        while not websocket.closed:
-            message = await websocket.recv()
+    while True:
+        try:
+            message = await websocket.receive()
             msg = json.loads(message)
             if "peer_group_id" in msg:
                 peer_group_id = msg["peer_group_id"]
                 if peer_group_id not in peer_list:
-                    peer_list[peer_group_id] = set() # Create new set for all peers in peer_group_id
-                peer_list[peer_group_id].add(websocket) # Add peer's socket to peer_group
+                    peer_list[peer_group_id] = set()
+                peer_list[peer_group_id].add(websocket._get_current_object())
             for peer in peer_list[peer_group_id]:
-                if peer != websocket:
+                if peer != websocket._get_current_object():
                     await peer.send(message)
-    except websockets.exceptions.ConnectionClosed as e:
-        pass
-    finally:
-        peer_list[peer_group_id].remove(websocket)
-        if len(peer_list[peer_group_id]) < 1:
-            peer_list.pop(peer_group_id)
+        except asyncio.exceptions.CancelledError:
+            peer_list[peer_group_id].remove(websocket._get_current_object())
+            if len(peer_list[peer_group_id]) < 1:
+                peer_list.pop(peer_group_id)
 
-async def main():
-    async with websockets.serve(
-        handler, "", 8001, ping_timeout=30,
-        logger=LoggerAdapter(logging.getLogger("websockets.server"), None),
-        ):
-        await asyncio.Future()  # run forever
-
-if __name__ == "__main__":
-    asyncio.run(main())
+app.run()