Files
virtuality/scripts/patch_web_console.py
T

102 строки
7.5 KiB
Python

#!/usr/bin/env python3
from pathlib import Path
import sys
app_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path('/opt/virtuality/web/app.py')
if not app_path.exists():
raise SystemExit(f'app.py not found: {app_path}')
text = app_path.read_text()
safe_proxy = '''# Virtuality noVNC console patch
async def proxy_vnc_to_websocket(reader: asyncio.StreamReader, websocket: WebSocket) -> None:
while True:
data = await reader.read(65536)
if not data:
break
try:
await websocket.send_bytes(data)
except (WebSocketDisconnect, RuntimeError, ConnectionError):
break
async def proxy_websocket_to_vnc(websocket: WebSocket, writer: asyncio.StreamWriter) -> None:
while True:
try:
message = await websocket.receive()
except (WebSocketDisconnect, RuntimeError, ConnectionError):
break
if message.get("type") == "websocket.disconnect":
break
if message.get("bytes") is not None:
writer.write(message["bytes"])
elif message.get("text") is not None:
writer.write(message["text"].encode())
try:
await writer.drain()
except (RuntimeError, ConnectionError, BrokenPipeError):
break
'''
safe_gather = ''' tasks = [
asyncio.create_task(proxy_vnc_to_websocket(reader, websocket)),
asyncio.create_task(proxy_websocket_to_vnc(websocket, writer)),
]
try:
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
for task in pending:
task.cancel()
await asyncio.gather(*pending, return_exceptions=True)
await asyncio.gather(*done, return_exceptions=True)
finally:
writer.close()
try:
await writer.wait_closed()
except Exception:
pass
'''
if 'Virtuality noVNC console patch' in text:
start = text.find('# Virtuality noVNC console patch')
marker = '\n\n@app.get("/login", response_class=HTMLResponse)'
end = text.find(marker, start)
if start != -1 and end != -1:
text = text[:start] + safe_proxy + text[end:]
old_gather = ''' try:
await asyncio.gather(proxy_vnc_to_websocket(reader, websocket), proxy_websocket_to_vnc(websocket, writer))
except (WebSocketDisconnect, asyncio.CancelledError, ConnectionError):
pass
finally:
writer.close()
await writer.wait_closed()
'''
if old_gather in text:
text = text.replace(old_gather, safe_gather, 1)
app_path.write_text(text)
print(f'noVNC console patch upgraded: {app_path}')
raise SystemExit(0)
text = text.replace('import crypt\n', 'import asyncio\nimport crypt\n', 1)
text = text.replace('from fastapi import FastAPI, Request, Form, UploadFile, File\n', 'from fastapi import FastAPI, Request, Form, UploadFile, File, WebSocket, WebSocketDisconnect\n', 1)
text = text.replace('serializer = URLSafeSerializer(SESSION_SECRET, salt="virtuality-session")\n', 'serializer = URLSafeSerializer(SESSION_SECRET, salt="virtuality-session")\nconsole_serializer = URLSafeSerializer(SESSION_SECRET, salt="virtuality-console")\n', 1)
text = text.replace('DEFAULT_BRIDGE = "br0"\n', 'DEFAULT_BRIDGE = "br0"\nNOVNC_DIR = next((p for p in [Path("/usr/share/novnc"), Path("/usr/share/novnc/app")] if p.exists()), None)\n', 1)
text = text.replace('app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")\n', 'app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")\nif NOVNC_DIR:\n app.mount("/novnc", StaticFiles(directory=str(NOVNC_DIR)), name="novnc")\n', 1)
old_get_user = '''def get_current_user(request: Request) -> str | None:\n token = request.cookies.get("virtuality_session")\n if not token:\n return None\n try:\n data = serializer.loads(token)\n except BadSignature:\n return None\n return AUTH_USER if data.get("user") == AUTH_USER else None\n'''
new_get_user = '''def user_from_session_token(token: str | None) -> str | None:\n if not token:\n return None\n try:\n data = serializer.loads(token)\n except BadSignature:\n return None\n return AUTH_USER if data.get("user") == AUTH_USER else None\n\n\ndef get_current_user(request: Request) -> str | None:\n return user_from_session_token(request.cookies.get("virtuality_session"))\n'''
text = text.replace(old_get_user, new_get_user, 1)
insert_after_vm_ip = '''def vm_vnc_display(name: str) -> str:\n return run_cmd(["virsh", "vncdisplay", name], timeout=8)["stdout"] or "not available"\n\n\ndef vnc_display_to_port(display: str) -> int | None:\n value = (display or "").strip()\n if not value or value == "not available":\n return None\n match = re.search(r":(\\d+)$", value)\n if not match:\n return None\n display_number = int(match.group(1))\n if display_number >= 5900:\n return display_number\n return 5900 + display_number\n\n\ndef console_info(name: str) -> dict[str, Any]:\n display = vm_vnc_display(name)\n port = vnc_display_to_port(display)\n has_novnc = bool(NOVNC_DIR and (NOVNC_DIR / "vnc.html").exists())\n token = None\n url = None\n if port and has_novnc:\n token = console_serializer.dumps({"vm": name, "port": port})\n url = f"/novnc/vnc.html?autoconnect=1&resize=scale&path=console/ws/{token}"\n return {"vm": name, "display": display, "port": port, "has_novnc": has_novnc, "novnc_dir": str(NOVNC_DIR) if NOVNC_DIR else "not installed", "url": url}\n\n\n'''
text = text.replace('def vm_details(name: str) -> dict[str, Any]:\n', insert_after_vm_ip + 'def vm_details(name: str) -> dict[str, Any]:\n', 1)
text = text.replace('"vnc": run_cmd(["virsh", "vncdisplay", name], timeout=8)["stdout"] or "not available",', '"vnc": vm_vnc_display(name),', 1)
text = text.replace('@app.get("/login", response_class=HTMLResponse)\n', safe_proxy + '\n@app.get("/login", response_class=HTMLResponse)\n', 1)
console_routes = '''\n\n@app.get("/vm/{name}/console", response_class=HTMLResponse)\ndef vm_console_page(request: Request, name: str):\n auth_redirect = require_auth(request)\n if auth_redirect:\n return auth_redirect\n if not valid_vm_name(name) or not vm_exists(name):\n return RedirectResponse(url="/", status_code=303)\n return templates.TemplateResponse("console.html", {"request": request, "app_name": APP_NAME, "user": AUTH_USER, "vm": vm_details(name), "console": console_info(name)})\n\n\n@app.websocket("/console/ws/{token}")\nasync def console_websocket(websocket: WebSocket, token: str):\n if user_from_session_token(websocket.cookies.get("virtuality_session")) != AUTH_USER:\n await websocket.close(code=1008)\n return\n try:\n payload = console_serializer.loads(token)\n vm_name = payload.get("vm")\n target_port = int(payload.get("port"))\n except Exception:\n await websocket.close(code=1008)\n return\n if not valid_vm_name(vm_name) or not vm_exists(vm_name) or target_port < 5900 or target_port > 5999:\n await websocket.close(code=1008)\n return\n await websocket.accept()\n try:\n reader, writer = await asyncio.open_connection("127.0.0.1", target_port)\n except Exception:\n await websocket.close(code=1011)\n return\n''' + safe_gather.replace('\n', '\n')
text = text.replace('\n\n@app.get("/vm/{name}", response_class=HTMLResponse)\n', console_routes + '\n\n@app.get("/vm/{name}", response_class=HTMLResponse)\n', 1)
app_path.write_text(text)
print(f'noVNC console patch applied: {app_path}')