170 lines
7.7 KiB
Python
170 lines
7.7 KiB
Python
"""WebSocket-прокси для VNC-консоли.
|
||
|
||
Поток данных:
|
||
1. Backend получает от Proxmox одноразовый тикет через `vncproxy` (REST API,
|
||
работает с API-токеном).
|
||
2. Открывает WebSocket к Proxmox с этим тикетом и портом. Аутентификация:
|
||
- API-токен (заголовок `Authorization: PVEAPIToken=...`) — новые PVE ≥ 7.x
|
||
- Cookie PVEAuthCookie — если заданы PVE_USERNAME + PVE_PASSWORD
|
||
3. Проксирует бинарный VNC-поток между браузером клиента и Proxmox.
|
||
|
||
Особенности Proxmox:
|
||
- В Proxmox 8.x+ vncticket может содержать спецсимволы (#, /, %, :),
|
||
которые ломают парсинг URL (RFC 3986). Кодируем их через quote().
|
||
- PVE может вернуть HTTP redirect (302) на `https://...` вместо `wss://...`;
|
||
перехватываем и переписываем scheme через process_redirect.
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
from typing import Optional, Tuple
|
||
from urllib.parse import quote
|
||
|
||
import httpx
|
||
import websockets
|
||
from websockets.http11 import Response
|
||
from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
|
||
|
||
from ..config import settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
router = APIRouter(prefix="/console", tags=["console"])
|
||
|
||
|
||
def _process_redirect(request, response: Response) -> Optional[str]:
|
||
"""Перехватывает HTTP-redirect от Proxmox и переписывает https://→wss://.
|
||
|
||
Proxmox иногда возвращает Location со схемой https:// (хотя для WebSocket
|
||
положено wss://). Библиотека websockets в этом случае падает с
|
||
InvalidURI("scheme isn't ws or wss"). Возвращаем исправленный URL.
|
||
"""
|
||
location = response.headers.get("Location", "")
|
||
if not location:
|
||
return None
|
||
if location.startswith("https://"):
|
||
return "wss://" + location[len("https://"):]
|
||
if location.startswith("http://"):
|
||
return "ws://" + location[len("http://"):]
|
||
return location
|
||
|
||
|
||
async def _get_pve_auth_cookie() -> Optional[Tuple[str, str]]:
|
||
"""Аутентифицируется в Proxmox по логину/паролю и возвращает (cookie_header, csrf_token).
|
||
|
||
Возвращает None, если PVE_USERNAME/PVE_PASSWORD не заданы в настройках.
|
||
"""
|
||
if not settings.pve_username or not settings.pve_password:
|
||
return None
|
||
|
||
url = f"{settings.pve_host}/api2/json/access/ticket"
|
||
payload = {"username": settings.pve_username, "password": settings.pve_password}
|
||
verify = settings.pve_verify_ssl
|
||
|
||
try:
|
||
async with httpx.AsyncClient(verify=verify, timeout=10.0) as client:
|
||
resp = await client.post(url, data=payload)
|
||
resp.raise_for_status()
|
||
data = resp.json().get("data", {})
|
||
cookie = data.get("ticket") # в Proxmox поле называется "ticket"
|
||
csrf = data.get("CSRFPreventionToken")
|
||
if not cookie:
|
||
logger.error("auth: Proxmox не вернул PVEAuthCookie (пустой ответ)")
|
||
return None
|
||
# Для WebSocket-эндпоинта vncwebsocket достаточно PVEAuthCookie.
|
||
cookie_header = f"PVEAuthCookie={cookie}"
|
||
logger.info("auth: PVEAuthCookie получен успешно")
|
||
return cookie_header, csrf or ""
|
||
except Exception as exc:
|
||
logger.exception("auth: ошибка аутентификации в Proxmox")
|
||
return None
|
||
|
||
|
||
@router.websocket("/ws")
|
||
async def console_ws(
|
||
websocket: WebSocket,
|
||
node: str = Query(...),
|
||
vmid: int = Query(...),
|
||
guest_type: str = Query(...),
|
||
port: int = Query(...),
|
||
ticket: str = Query(...),
|
||
):
|
||
"""Проксирует бинарный VNC-поток между браузером и Proxmox."""
|
||
await websocket.accept()
|
||
|
||
guest_path = "qemu" if guest_type == "vm" else "lxc"
|
||
pve_host_only = settings.pve_host.replace("http://", "").replace("https://", "")
|
||
# quote() экранирует спецсимволы в тикете и порте по правилам RFC 3986,
|
||
# не затрагивая уже экранированные символы. Без этого в Proxmox 8.x+
|
||
# vncticket может содержать # / % : которые ломают URL.
|
||
safe_ticket = quote(ticket, safe="")
|
||
safe_port = quote(str(port), safe="")
|
||
upstream_url = (
|
||
f"wss://{pve_host_only}/api2/json/nodes/{node}/{guest_path}/{vmid}/vncwebsocket"
|
||
f"?port={safe_port}&vncticket={safe_ticket}"
|
||
)
|
||
|
||
# Определяем способ аутентификации.
|
||
# Cookie-вариант приоритетнее, т.к. старые/строгие PVE не принимают API-токен на WS.
|
||
headers: list = []
|
||
auth = await _get_pve_auth_cookie()
|
||
if auth:
|
||
cookie_header, _ = auth
|
||
headers.append(("Cookie", cookie_header))
|
||
auth_kind = "cookie"
|
||
else:
|
||
headers.append(("Authorization", f"PVEAPIToken={settings.pve_token_name}={settings.pve_token_value}"))
|
||
auth_kind = "api-token"
|
||
|
||
logger.info(
|
||
"console_ws: connecting node=%s vmid=%s guest=%s port=%s auth=%s",
|
||
node, vmid, guest_type, port, auth_kind,
|
||
)
|
||
|
||
try:
|
||
# extra_headers (а не additional_headers) — последний не поддерживается
|
||
# в uvloop, который использует uvicorn[standard].
|
||
async with websockets.connect(
|
||
upstream_url,
|
||
extra_headers=headers,
|
||
subprotocols=["binary"],
|
||
ssl=None if settings.pve_verify_ssl else False,
|
||
process_redirect=_process_redirect,
|
||
) as upstream:
|
||
logger.info("console_ws: upstream connected")
|
||
|
||
async def client_to_upstream():
|
||
try:
|
||
while True:
|
||
data = await websocket.receive_bytes()
|
||
await upstream.send(data)
|
||
except WebSocketDisconnect:
|
||
logger.info("console_ws: client disconnected")
|
||
except Exception as exc:
|
||
logger.warning("console_ws: client_to_upstream error: %s", exc)
|
||
|
||
async def upstream_to_client():
|
||
try:
|
||
async for message in upstream:
|
||
if isinstance(message, str):
|
||
message = message.encode()
|
||
await websocket.send_bytes(message)
|
||
except websockets.ConnectionClosed as exc:
|
||
logger.info("console_ws: upstream closed: code=%s", exc.code)
|
||
except Exception as exc:
|
||
logger.warning("console_ws: upstream_to_client error: %s", exc)
|
||
|
||
await asyncio.gather(client_to_upstream(), upstream_to_client())
|
||
except websockets.exceptions.InvalidStatus as exc:
|
||
status = exc.response.status_code if exc.response else "?"
|
||
logger.error(
|
||
"console_ws: Proxmox отклонил подключение (status=%s, auth=%s). "
|
||
"Если 401 — задайте PVE_USERNAME/PVE_PASSWORD в backend/.env "
|
||
"для cookie-аутентификации.",
|
||
status, auth_kind,
|
||
)
|
||
await websocket.close(code=1011, reason=f"Proxmox auth failed: {status}")
|
||
except Exception as exc: # noqa: BLE001
|
||
logger.exception("console_ws: ошибка подключения к Proxmox")
|
||
await websocket.close(code=1011, reason=str(exc)[:120])
|