81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
import aiohttp
|
|
from aiohttp import web
|
|
import asyncio
|
|
from datetime import datetime
|
|
|
|
minirt_vnc_host = "localhost"
|
|
minirt_vnc_port = 7080
|
|
|
|
connections = {}
|
|
max_duration = 60
|
|
|
|
|
|
async def proxy_http(request):
|
|
service_name = request.match_info.get('service')
|
|
target = request.match_info.get('target')
|
|
|
|
if service_name == 'minirt':
|
|
target = f"{minirt_vnc_host}:{minirt_vnc_port}/{target}"
|
|
elif service_name == 'minishell':
|
|
target = ''
|
|
|
|
print(f"proxying to host {target}")
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.request(
|
|
request.method,
|
|
f"http://{target}",
|
|
headers=request.headers,
|
|
data=await request.read()
|
|
) as resp:
|
|
return web.Response(
|
|
body=await resp.read(),
|
|
status=resp.status,
|
|
headers=resp.headers
|
|
)
|
|
|
|
async def proxy_ws(request: web.Request):
|
|
service_name = request.match_info.get('service')
|
|
|
|
if service_name in connections:
|
|
return web.Response(text='Service already in use', status=409)
|
|
|
|
ws = web.WebSocketResponse()
|
|
await ws.prepare(request)
|
|
connections[service_name] = (ws, datetime.now())
|
|
|
|
target = ''
|
|
if service_name == 'minirt':
|
|
target = f"{minirt_vnc_host}:{minirt_vnc_port}"
|
|
elif service_name == 'minishell':
|
|
target = ''
|
|
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.ws_connect(f"ws://{target}/websockify") as remote_ws:
|
|
async def forward_to_remote():
|
|
async for msg in ws:
|
|
if msg.type == aiohttp.WSMsgType.BINARY:
|
|
await remote_ws.send_bytes(msg.data)
|
|
|
|
async def forward_to_client():
|
|
async for msg in remote_ws:
|
|
if msg.type == aiohttp.WSMsgType.BINARY:
|
|
await ws.send_bytes(msg.data)
|
|
|
|
async def check_timeout():
|
|
while True:
|
|
await asyncio.sleep(10)
|
|
elapsed = (datetime.now() - connections[service_name][1]).total_seconds()
|
|
if elapsed > max_duration:
|
|
await ws.close()
|
|
break
|
|
|
|
await asyncio.gather(forward_to_remote(), forward_to_client(), check_timeout())
|
|
finally:
|
|
if service_name in connections:
|
|
del connections[service_name]
|
|
|
|
return ws
|
|
|
|
|