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
socket_id = -1
async def handle_client(client_reader, client_writer):
async def handle_client(socket_reader, socket_writer):
global socket_id
socket_id += 1
try:
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
async def client_to_websocket():
while True:
data = await client_reader.read(2024)
data = await socket_reader.read(2024)
print(f'TCP{socket_id}>WS: ', hashlib.md5(data).hexdigest())
if not data:
break
c = conn.Conn(socket_id, data)
await websocket.send(c.to_ws_bytes())
c = conn.WSMsg(socket_id, conn.MsgType.DATA, data)
await websocket.send(c.to_bytes())
# Forwarding data from WebSocket to client
async def websocket_to_client():
while True:
message = await websocket.recv()
c = conn.Conn.from_ws_bytes(message)
if c.socketid == socket_id:
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.msg == conn.MsgType.DISCONNECT:
print(f"Client {socket_id} disconnected")
break
else:
print(f'WS>TCP@{socket_id}: ', hashlib.md5(message).hexdigest())
client_writer.write(c.data)
await client_writer.drain()
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
await asyncio.gather(client_to_websocket(), websocket_to_client())
print(f">>>> Client {socket_id} disconnected")
except Exception as e:
print(f"Error: {e}")
import traceback
traceback.print_exc()
finally:
client_writer.close()
await client_writer.wait_closed()
print(f">>>>>>>>>> Closing client {socket_id}")
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():
port = PORT

40
conn.py
View file

@ -1,17 +1,41 @@
from typing import Optional
import bencoder
from enum import Enum
class Conn:
class MsgType(Enum):
CONNECT = 0
DISCONNECT = 1
DATA = 2
class WSMsg:
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.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
def from_ws_bytes(b: bytes) -> 'Conn':
def from_bytes(b: bytes) -> 'WSMsg':
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)

76
pod.py
View file

@ -9,60 +9,24 @@ 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
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):
try:
print(f"New socket: {socketid}. Waiting on recv")
print(tcpreader)
while True:
data = await tcpreader.read(2024)
print('GOT DATA', data)
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
print(f'TCP@{socketid}>WS: RAW', data)
c = conn.Conn(socketid, data)
c = conn.WSMsg(socketid, conn.MsgType.DATA, data)
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:
print(f"Error: {e}")
print(f"{socketid} Error: {e}")
import traceback
traceback.print_exc()
@ -70,22 +34,34 @@ async def handle_socket_read(socketid, tcpreader, ws):
async def handle_ws_incoming(ws, sockets):
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
if socketid not in sockets:
if c.msg == conn.MsgType.CONNECT:
if socketid in sockets:
print(f"Socket {socketid} already connected")
return
else:
print(f"New socket: {socketid}")
tcpreader, tcpwriter = await asyncio.open_connection(HOST, PORT)
sockets[socketid] = (tcpreader, tcpwriter)
print(f'TCPR: {tcpreader}')
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]
print(f'WS@{socketid}>TCP: ', hashlib.md5(data).hexdigest())
print(tcpwriter)
tcpwriter.write(c.data)
tcpwriter.write(c.payload)