diff --git a/Pipfile b/Pipfile index b17b6d3..ca5186a 100644 --- a/Pipfile +++ b/Pipfile @@ -6,6 +6,7 @@ name = "pypi" [packages] websockets = "*" asyncio = "*" +bencoder = "*" [dev-packages] diff --git a/Pipfile.lock b/Pipfile.lock index 90cfae8..1b25a6e 100644 --- a/Pipfile.lock +++ b/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", diff --git a/client.py b/client.py index 4543d84..ca998a0 100644 --- a/client.py +++ b/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() diff --git a/conn.py b/conn.py new file mode 100644 index 0000000..bb333a5 --- /dev/null +++ b/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"]) diff --git a/pod.py b/pod.py index 0523557..78d6e9d 100644 --- a/pod.py +++ b/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) diff --git a/proxy.py b/proxy.py index 3af31f3..9ecd140 100644 --- a/proxy.py +++ b/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()