Init

transphase.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+import asyncio
+import websockets
+import logging 
+import json
+
+
+logging.basicConfig(
+    format="%(asctime)s %(message)s",
+    level=logging.INFO,
+)
+
+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
+
+async def handler(websocket):
+    peer_group_id = None
+    try:
+        while not websocket.closed:
+            message = await websocket.recv()
+            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
+            for peer in peer_list[peer_group_id]:
+                if peer != websocket:
+                    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)
+
+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())