this shit fucking works

Signed-off-by: Frank Villaro-Dixon <frank@villaro-dixon.eu>
This commit is contained in:
Frank Villaro-Dixon 2024-08-27 23:21:22 +02:00
parent c80f43b5f5
commit 80184c4b55
3 changed files with 96 additions and 77 deletions

View file

@ -13,43 +13,62 @@ WEBSOCKET_URL = 'ws://localhost:9999/data'
ALREADY_DONE = False ALREADY_DONE = False
socket_id = -1 socket_id = -1
async def handle_client(client_reader, client_writer): async def handle_client(socket_reader, socket_writer):
global socket_id global socket_id
socket_id += 1 socket_id += 1
try: try:
async with websockets.connect(WEBSOCKET_URL) as websocket: async with websockets.connect(WEBSOCKET_URL) as websocket:
print(f"New client connected socket {socket_id}: {socket_reader} {socket_writer}")
m = conn.WSMsg(socket_id, conn.MsgType.CONNECT)
await websocket.send(m.to_bytes())
# Forwarding data from client to WebSocket # Forwarding data from client to WebSocket
async def client_to_websocket(): async def client_to_websocket():
while True: while True:
data = await client_reader.read(2024) data = await socket_reader.read(2024)
print(f'TCP{socket_id}>WS: ', hashlib.md5(data).hexdigest()) print(f'TCP{socket_id}>WS: ', hashlib.md5(data).hexdigest())
if not data: if not data:
break break
c = conn.Conn(socket_id, data) c = conn.WSMsg(socket_id, conn.MsgType.DATA, data)
await websocket.send(c.to_ws_bytes()) await websocket.send(c.to_bytes())
# Forwarding data from WebSocket to client # Forwarding data from WebSocket to client
async def websocket_to_client(): async def websocket_to_client():
while True: while True:
message = await websocket.recv() message = await websocket.recv()
c = conn.Conn.from_ws_bytes(message) c = conn.WSMsg.from_bytes(message)
# XXX this is ugly, because it means that the data is sent twice or more if 2+ connections..
if c.socketid == socket_id: if c.socketid == socket_id:
# XXX this is ugly, because it means that the data is sent twice or more if 2+ connections.. if c.msg == conn.MsgType.DISCONNECT:
print(f'WS>TCP@{socket_id}: ', hashlib.md5(message).hexdigest()) print(f"Client {socket_id} disconnected")
client_writer.write(c.data) break
await client_writer.drain() else:
print(f'WS>TCP@{socket_id}: ', hashlib.md5(message).hexdigest())
socket_writer.write(c.payload)
await socket_writer.drain()
else:
print(f'WS>TCP@{socket_id}: ', hashlib.md5(message).hexdigest(), 'skipping')
# Run both tasks concurrently # Run both tasks concurrently
await asyncio.gather(client_to_websocket(), websocket_to_client()) await asyncio.gather(client_to_websocket(), websocket_to_client())
print(f">>>> Client {socket_id} disconnected")
except Exception as e: except Exception as e:
print(f"Error: {e}") print(f"Error: {e}")
import traceback import traceback
traceback.print_exc() traceback.print_exc()
finally: finally:
client_writer.close() print(f">>>>>>>>>> Closing client {socket_id}")
await client_writer.wait_closed() m = conn.WSMsg(socket_id, conn.MsgType.DISCONNECT)
await websocket.send(m.to_bytes())
socket_writer.close()
await socket_writer.wait_closed()
async def main(): async def main():
port = PORT port = PORT

40
conn.py
View file

@ -1,17 +1,41 @@
from typing import Optional
import bencoder import bencoder
from enum import Enum
class Conn:
class MsgType(Enum):
CONNECT = 0
DISCONNECT = 1
DATA = 2
class WSMsg:
socketid: int socketid: int
data: bytes msg: MsgType
payload: Optional[bytes]
def __init__(self, socketid: int, data: bytes): def __init__(self, socketid: int, msg: MsgType, payload: Optional[bytes] = None):
self.socketid = socketid self.socketid = socketid
self.data = data self.msg = msg
self.payload = payload
def to_ws_bytes(self) -> bytes:
return bencoder.encode({b"socketid": self.socketid, b"data": self.data})
@staticmethod @staticmethod
def from_ws_bytes(b: bytes) -> 'Conn': def from_bytes(b: bytes) -> 'WSMsg':
d = bencoder.decode(b) d = bencoder.decode(b)
return Conn(d[b"socketid"], d[b"data"]) socketid = d[b"sid"]
mtype = MsgType(d[b"mt"])
if mtype == MsgType.CONNECT:
return WSMsg(socketid, mtype)
elif mtype == MsgType.DISCONNECT:
return WSMsg(socketid, mtype)
elif mtype == MsgType.DATA:
return WSMsg(socketid, mtype, d[b"data"])
def to_bytes(self) -> bytes:
d = {b"sid": self.socketid, b"mt": self.msg.value}
if self.payload:
d[b"data"] = self.payload
return bencoder.encode(d)

88
pod.py
View file

@ -9,60 +9,24 @@ HOST = '192.168.21.30'
PORT = 6443 PORT = 6443
WEBSOCKET_URL = 'ws://localhost:9999/data' 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
async def tcp_to_websocket():
print("tcp_to_websocket...")
while True:
data = await tcpreader.read(2024)
c = conn.Conn(0, data)
socketid = c.socketid
print(f'TCP@{socketid}>WS: ', hashlib.md5(data).hexdigest())
if not data:
break
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()
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...")
# Run both tasks concurrently
await asyncio.gather(tcp_to_websocket(), websocket_to_tcp())
print('done')
except Exception as e:
print(f"ASYNCIOD Error: {e}")
async def handle_socket_read(socketid, tcpreader, ws): async def handle_socket_read(socketid, tcpreader, ws):
try: try:
print(f"New socket: {socketid}. Waiting on recv") print(f"New socket: {socketid}. Waiting on recv")
print(tcpreader)
while True: while True:
data = await tcpreader.read(2024) data = await tcpreader.read(2024)
print('GOT DATA', data)
if data == b'': if data == b'':
print(f"Connection closed: {socketid}") print(f"TCP@{socketid} Connection closed")
c = conn.WSMsg(socketid, conn.MsgType.DISCONNECT)
await ws.send(c.to_bytes())
break break
print(f'TCP@{socketid}>WS: RAW', data) c = conn.WSMsg(socketid, conn.MsgType.DATA, data)
c = conn.Conn(socketid, data)
print(f'TCP@{socketid}>WS: ', hashlib.md5(data).hexdigest()) print(f'TCP@{socketid}>WS: ', hashlib.md5(data).hexdigest())
await ws.send(c.to_ws_bytes()) await ws.send(c.to_bytes())
except Exception as e: except Exception as e:
print(f"Error: {e}") print(f"{socketid} Error: {e}")
import traceback import traceback
traceback.print_exc() traceback.print_exc()
@ -70,22 +34,34 @@ async def handle_socket_read(socketid, tcpreader, ws):
async def handle_ws_incoming(ws, sockets): async def handle_ws_incoming(ws, sockets):
data = await ws.recv() data = await ws.recv()
c = conn.Conn.from_ws_bytes(data) c = conn.WSMsg.from_bytes(data)
print(f'NEW DATA: {c.socketid} {c.msg} {c.payload}')
socketid = c.socketid socketid = c.socketid
if socketid not in sockets: if c.msg == conn.MsgType.CONNECT:
print(f"New socket: {socketid}") if socketid in sockets:
tcpreader, tcpwriter = await asyncio.open_connection(HOST, PORT) print(f"Socket {socketid} already connected")
sockets[socketid] = (tcpreader, tcpwriter) return
print(f'TCPR: {tcpreader}') else:
asyncio.create_task(handle_socket_read(socketid, tcpreader, ws)) print(f"New socket: {socketid}")
else: tcpreader, tcpwriter = await asyncio.open_connection(HOST, PORT)
sockets[socketid] = (tcpreader, tcpwriter)
asyncio.create_task(handle_socket_read(socketid, tcpreader, ws))
elif c.msg == conn.MsgType.DISCONNECT:
if socketid not in sockets:
print(f"Socket {socketid} not connected")
return
else:
print(f"Socket {socketid} disconnected")
del sockets[socketid]
tcpwriter.close()
await tcpwriter.wait_closed()
elif c.msg == conn.MsgType.DATA:
tcpreader, tcpwriter = sockets[socketid] tcpreader, tcpwriter = sockets[socketid]
print(f'WS@{socketid}>TCP: ', hashlib.md5(data).hexdigest())
print(f'WS@{socketid}>TCP: ', hashlib.md5(data).hexdigest()) tcpwriter.write(c.payload)
print(tcpwriter)
tcpwriter.write(c.data)