88 lines
2.5 KiB
Python
88 lines
2.5 KiB
Python
import time
|
|
import asyncio
|
|
from asyncio.protocols import DatagramProtocol
|
|
|
|
import websockets
|
|
|
|
TX_HOST = '127.0.0.1'
|
|
TX_PORT = '80'
|
|
TX_ROUTE = '/'
|
|
RX_HOST = '0.0.0.0'
|
|
RX_PORT = 5005
|
|
|
|
class MyProto(DatagramProtocol):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def connection_made(self, transport):
|
|
self.transport = transport
|
|
self.tunnels = {}
|
|
self.last_packet_time = {}
|
|
|
|
async def ws_to_udp_loop(self, first_packet, addr):
|
|
try:
|
|
ws = await websockets.connect("ws://"+TX_HOST+':'+TX_PORT+TX_ROUTE)
|
|
self.tunnels[addr] = ws
|
|
print("WebSocket for", addr, "created")
|
|
except Exception as e:
|
|
print("Unable to connect, check if remote server is available. Exception text:", e)
|
|
return
|
|
|
|
try:
|
|
while True:
|
|
await ws.send(first_packet)
|
|
async for message in ws:
|
|
self.transport.sendto(message, addr)
|
|
self.last_packet_time[addr] = time.time()
|
|
except Exception as e:
|
|
print("Received exception on WebSocket for", addr,"-", e)
|
|
finally:
|
|
print("WebSocket for", addr, "closed")
|
|
self.tunnels.pop(addr, None)
|
|
self.last_packet_time.pop(addr, None)
|
|
await ws.close()
|
|
|
|
|
|
def datagram_received(self, data, addr):
|
|
loop = asyncio.get_running_loop()
|
|
ws = self.tunnels.get(addr, None)
|
|
|
|
if ws is None:
|
|
loop.create_task(self.ws_to_udp_loop(data, addr))
|
|
else:
|
|
loop.create_task(ws.send(data))
|
|
self.last_packet_time[addr] = time.time()
|
|
|
|
async def cleanup(self):
|
|
connections_for_removal = []
|
|
for addr in self.last_packet_time.keys():
|
|
time_passed = time.time() - self.last_packet_time[addr]
|
|
if time_passed > 3600:
|
|
print(f"Closing connection for {addr} (inactive for: {int(time_passed)} seconds)")
|
|
connections_for_removal.append(addr)
|
|
|
|
for addr in connections_for_removal:
|
|
await self.tunnels[addr].close()
|
|
|
|
|
|
|
|
async def main():
|
|
addr = RX_HOST
|
|
proto = MyProto()
|
|
transport, protocol = await asyncio.get_event_loop().create_datagram_endpoint(
|
|
lambda: proto,
|
|
local_addr=(addr, RX_PORT)
|
|
)
|
|
|
|
print(f'Serving on {addr}:{RX_PORT}')
|
|
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(3600)
|
|
await proto.cleanup()
|
|
|
|
finally:
|
|
transport.close()
|
|
|
|
asyncio.run(main())
|