84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
import aiohttp
|
|
from aiohttp import web
|
|
import asyncio
|
|
from datetime import datetime
|
|
|
|
from services import services
|
|
|
|
|
|
connections = {}
|
|
max_duration = 60
|
|
|
|
async def proxy_http(request: web.Request):
|
|
service_name = request.match_info.get('service')
|
|
|
|
if not service_name or not services[service_name]:
|
|
return web.json_response(status=404, data={ "success": "false", "message": "Unknown service name" })
|
|
url = services[service_name]
|
|
|
|
target = request.match_info.get('target')
|
|
target = f"{url}/{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)
|
|
|
|
if not service_name or not services[service_name]:
|
|
return web.json_response(status=404, data={ "success": "false", "message": "Unknown service name" })
|
|
|
|
ws = web.WebSocketResponse()
|
|
await ws.prepare(request)
|
|
|
|
connections[service_name] = (ws, datetime.now())
|
|
|
|
url = services[service_name]
|
|
|
|
target = request.match_info.get('target')
|
|
target = f"{url}/{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
|
|
|
|
|