feat: Some quality of life improvements

proxy/websocket_client.py

@@ -0,0 +1,85 @@
+import asyncio
+import logging
+import websockets
+
+from config import TARGET_WS_URL, LOCAL_HOST, LOCAL_PORT
+from selenium_utils import fetch_fresh_credentials
+
+log = logging.getLogger(__name__)
+
+
+async def handle_connection(local_client):
+    log.info("Local client connected, fetching credentials...")
+
+    result = await asyncio.get_running_loop().run_in_executor(
+        None, fetch_fresh_credentials
+    )
+    if result is None:
+        log.error("Could not fetch credentials")
+        await local_client.close(1011, "Failed to bypass Cloudflare")
+        return
+
+    user_agent, cookies = result
+    headers = {
+        "User-Agent": user_agent,
+        "Cookie": cookies,
+    }
+
+    try:
+        log.info("Connecting to OmegleWeb...")
+        async with websockets.connect(
+            TARGET_WS_URL,
+            additional_headers=headers,
+        ) as target_server:
+            log.info("OmegleWeb connected! Notifying client...")
+            await local_client.send("SUCCESS")
+            log.info("Starting relay")
+            await relay_bidirectional(local_client, target_server)
+
+    except websockets.exceptions.InvalidStatusCode as e:
+        log.error(f"Omegle connection failed (Status {e.status_code})")
+        await local_client.close(1011, f"Connection failed: {e.status_code}")
+
+    except websockets.exceptions.ConnectionClosed:
+        log.info("OmegleWeb disconnected")
+
+    except Exception as e:
+        log.error(f"Connection error: {type(e).__name__}: {e}")
+        await local_client.close(1011, str(e))
+
+
+async def relay_bidirectional(local_ws, target_ws):
+    async def forward_local():
+        try:
+            async for message in local_ws:
+                log.debug(f"[Local -> Omegle]: {message}")
+                await target_ws.send(message)
+        except websockets.exceptions.ConnectionClosed:
+            log.info("Local client disconnected")
+        except Exception as e:
+            log.error(f"Forward local error: {e}")
+
+    async def forward_remote():
+        try:
+            async for message in target_ws:
+                log.debug(f"[Omegle -> Local]: {message}")
+                await local_ws.send(message)
+        except websockets.exceptions.ConnectionClosed:
+            log.info("OmegleWeb disconnected")
+        except Exception as e:
+            log.error(f"Forward remote error: {e}")
+
+    await asyncio.gather(
+        forward_local(),
+        forward_remote(),
+    )
+
+
+async def main():
+    log.info(f"Starting WebSocket proxy on ws://{LOCAL_HOST}:{LOCAL_PORT}")
+    await websockets.serve(handle_connection, LOCAL_HOST, LOCAL_PORT)
+    await asyncio.Future()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())