pocv2

Signed-off-by: Frank Villaro-Dixon <frank@villaro-dixon.eu>

Pipfile

@@ -6,6 +6,7 @@ name = "pypi"
 [packages]
 websockets = "*"
 asyncio = "*"
+bencoder = "*"
 
 [dev-packages]
 

Pipfile.lock

@@ -1,7 +1,7 @@
 {
     "_meta": {
         "hash": {
-            "sha256": "2f522d4ca76f3a562a33486191d9a8f4fad53f9b7337675f8762a0a3bb880272"
+            "sha256": "32897fbdc87627e0a506eca3bb31e5b94fecf1d0b3ee5a3426c5def16fea1acc"
         },
         "pipfile-spec": 6,
         "requires": {
@@ -26,6 +26,13 @@
             "index": "pypi",
             "version": "==3.4.3"
         },
+        "bencoder": {
+            "hashes": [
+                "sha256:ac436f33fdd7e75b2d9057491d2abeefb5631efb13c813403db0adb11ab52fc2"
+            ],
+            "index": "pypi",
+            "version": "==0.2.0"
+        },
         "websockets": {
             "hashes": [
                 "sha256:02cc9bb1a887dac0e08bf657c5d00aa3fac0d03215d35a599130c2034ae6663a",

client.py

@@ -2,18 +2,20 @@ import asyncio
 import websockets
 import hashlib
 
+
+import conn
+
+
 LOCALHOST = '127.0.0.1'
 PORT = 8443
 WEBSOCKET_URL = 'ws://localhost:9999/data'
 
 ALREADY_DONE = False
+socket_id = -1
 
 async def handle_client(client_reader, client_writer):
-    global ALREADY_DONE
-    if ALREADY_DONE:
-        print('Will ret a 2nd time')
-        return
-    ALREADY_DONE = True
+    global socket_id
+    socket_id += 1
 
     try:
         async with websockets.connect(WEBSOCKET_URL) as websocket:
@@ -21,24 +23,30 @@ async def handle_client(client_reader, client_writer):
             async def client_to_websocket():
                 while True:
                     data = await client_reader.read(2024)
-                    print('TCP>WS: ', hashlib.md5(data).hexdigest())
+                    print(f'TCP{socket_id}>WS: ', hashlib.md5(data).hexdigest())
                     if not data:
                         break
-                    await websocket.send(data)
+                    c = conn.Conn(socket_id, data)
+                    await websocket.send(c.to_ws_bytes())
             
             # Forwarding data from WebSocket to client
             async def websocket_to_client():
                 while True:
                     message = await websocket.recv()
-                    print('WS>TCP: ', hashlib.md5(message).hexdigest())
-                    client_writer.write(message)
-                    await client_writer.drain()
+                    c = conn.Conn.from_ws_bytes(message)
+                    if c.socketid == socket_id:
+                        # XXX this is ugly, because it means that the data is sent twice or more if 2+ connections..
+                        print(f'WS>TCP@{socket_id}: ', hashlib.md5(message).hexdigest())
+                        client_writer.write(c.data)
+                        await client_writer.drain()
             
             # Run both tasks concurrently
             await asyncio.gather(client_to_websocket(), websocket_to_client())
 
     except Exception as e:
         print(f"Error: {e}")
+        import traceback
+        traceback.print_exc()
     finally:
         client_writer.close()
         await client_writer.wait_closed()

conn.py

@@ -0,0 +1,17 @@
+import bencoder
+
+class Conn:
+    socketid: int
+    data: bytes
+
+    def __init__(self, socketid: int, data: bytes):
+        self.socketid = socketid
+        self.data = data
+
+    def to_ws_bytes(self) -> bytes:
+        return bencoder.encode({b"socketid": self.socketid, b"data": self.data})
+
+    @staticmethod
+    def from_ws_bytes(b: bytes) -> 'Conn':
+        d = bencoder.decode(b)
+        return Conn(d[b"socketid"], d[b"data"])

pod.py

@@ -3,11 +3,15 @@ import socket
 import websockets
 import hashlib
 
+import conn
+
 HOST = '192.168.21.30'
 PORT = 6443
 WEBSOCKET_URL = 'ws://localhost:9999/data'
 
 async def handle_client(tcpreader, tcpwriter, ws):
+    sockets = {}
+
     print(f"New client connected: {tcpreader} {tcpwriter}")
     try:
         # Forwarding data from client to WebSocket
@@ -15,18 +19,22 @@ async def handle_client(tcpreader, tcpwriter, ws):
             print("tcp_to_websocket...")
             while True:
                 data = await tcpreader.read(2024)
-                print('TCP>WS: ', hashlib.md5(data).hexdigest())
+                c = conn.Conn(0, data)
+                socketid = c.socketid
+                print(f'TCP@{socketid}>WS: ', hashlib.md5(data).hexdigest())
                 if not data:
                     break
-                await ws.send(data)
+                await ws.send(c.to_ws_bytes())
         
         # Forwarding data from WebSocket to client
         async def websocket_to_tcp():
             print("websocket_to_tcp...")
             while True:
                 message = await ws.recv()
-                print('WS>TCP: ', hashlib.md5(message).hexdigest())
-                tcpwriter.write(message)
+                c = conn.Conn.from_ws_bytes(message)
+                socketid = c.socketid
+                print(f'WS>TCP@{socketid}: ', hashlib.md5(message).hexdigest())
+                tcpwriter.write(c.data)
                 await tcpwriter.drain()
 
         print("Running both tasks concurrently...")
@@ -38,8 +46,6 @@ async def handle_client(tcpreader, tcpwriter, ws):
         print(f"ASYNCIOD Error: {e}")
 
 async def main():
-    conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    conn.connect((HOST, PORT))
 
 
     tcpreader, tcpwriter = await asyncio.open_connection(HOST, PORT)

proxy.py

@@ -2,12 +2,15 @@ import asyncio
 import websockets
 import hashlib
 
+import websockets.asyncio.server
+
 # List to store connected clients
 connected_clients = set()
 
-async def handler(websocket, path):
+async def handler(websocket):
     # Register the new client
     print(f"New client connected: {websocket}")
+    print(f"WRP: {websocket.request.path}")
     connected_clients.add(websocket)
     try:
         async for message in websocket:
@@ -25,7 +28,7 @@ async def handler(websocket, path):
 
 async def main():
     # Start the WebSocket server
-    server = await websockets.serve(handler, "localhost", 9999)
+    server = await websockets.asyncio.server.serve(handler, "localhost", 9999)
     print("WebSocket server listening on ws://localhost:9999")
     await server.wait_closed()