website/proxy.py

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