127 lines
4.5 KiB
Python
127 lines
4.5 KiB
Python
# websocket_server/events/vnc.py
|
|
|
|
from flask_socketio import emit
|
|
import logging
|
|
import requests
|
|
from flask import request
|
|
|
|
from websocket_server.events import base
|
|
from websocket_server.shared_state import connected_workers
|
|
|
|
|
|
logger = logging.getLogger("websocket_server")
|
|
socketio = base.socketio
|
|
|
|
api_server_url = "http://172.17.0.1:5000/api"
|
|
|
|
# In-memory request map: vnc_request_id -> {worker_id, vnc_proxy_sid, vm_id}
|
|
vnc_request_worker_map = {}
|
|
def register_socketio_handlers(socketio):
|
|
|
|
|
|
|
|
|
|
@socketio.on("vnc_frame_from_worker")
|
|
def vnc_frame_from_worker(payload):
|
|
vnc_request_id = payload.get("vnc_request_id")
|
|
if not vnc_request_id:
|
|
logger.warning("Received VNC frame without request ID")
|
|
return
|
|
|
|
context = vnc_request_worker_map.get(vnc_request_id)
|
|
if not context:
|
|
logger.warning(f"No mapping found for VNC request {vnc_request_id}")
|
|
return
|
|
|
|
proxy_sid = context["vnc_proxy_sid"]
|
|
emit("vnc_frame_from_worker", payload, to=proxy_sid)
|
|
logger.debug(f"Forwarded VNC frame for {vnc_request_id} to VNC proxy SID {proxy_sid}")
|
|
|
|
|
|
@socketio.on("vnc_frame_from_novnc")
|
|
def vnc_frame_from_novnc(payload):
|
|
logger.debug(f"Got a frame from novnc")
|
|
#Send the VNC frame back to the VNC Proxy that matches the requestID
|
|
# logger.debug(f"Scored a VNC frame from novnc {data}")
|
|
emit("vnc_frame_from_novnc",payload,broadcast=True)
|
|
|
|
|
|
@socketio.on("start_vnc_stream")
|
|
def start_vnc_stream(data):
|
|
logger.info(f"Got VNC start stream request {data}")
|
|
vnc_request_id = data.get("vnc_request_id")
|
|
virtual_machine_id = data.get("virtual_machine_id")
|
|
user_token = data.get("user_token")
|
|
vnc_proxy_sid = request.sid # Who sent the request
|
|
logger.debug(f"User token: {user_token}")
|
|
if not vnc_request_id or not virtual_machine_id or not user_token:
|
|
logger.error("Missing required data in VNC stream request")
|
|
return
|
|
|
|
try:
|
|
vm_url = f"{api_server_url}/workloads/virtual_machines/{virtual_machine_id}"
|
|
response = requests.get(vm_url, timeout=5)
|
|
response.raise_for_status()
|
|
vm_info = response.json()
|
|
logger.debug(f"Fetched VM info: {vm_info}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch VM info for {virtual_machine_id}: {e}")
|
|
return
|
|
|
|
workload_host_id = vm_info.get("workload_host_id")
|
|
if not workload_host_id:
|
|
logger.error(f"No workload_host_id found for VM {virtual_machine_id}")
|
|
return
|
|
|
|
socket_id = connected_workers.get(workload_host_id)
|
|
if not socket_id:
|
|
logger.warning(f"Worker {workload_host_id} not connected")
|
|
return
|
|
|
|
# Save mapping with proxy SID
|
|
vnc_request_worker_map[vnc_request_id] = {
|
|
"worker_id": workload_host_id,
|
|
"vnc_proxy_sid": vnc_proxy_sid,
|
|
"virtual_machine_id": virtual_machine_id
|
|
}
|
|
|
|
emit("start_vnc_stream_on_worker", {
|
|
"vnc_request_id": vnc_request_id,
|
|
"virtual_machine_id": virtual_machine_id,
|
|
"workload_host_id": workload_host_id,
|
|
"vnc_port": 5900,
|
|
"vnc_ip_address": "127.0.0.1"
|
|
}, to=socket_id)
|
|
|
|
logger.info(f"Dispatched VNC stream request {vnc_request_id} to worker {workload_host_id}, from proxy {vnc_proxy_sid}")
|
|
|
|
@socketio.on("stop_vnc_stream")
|
|
def stop_vnc_stream(data):
|
|
logger.info(f"Got VNC stop stream request {data}")
|
|
vnc_request_id = data.get("vnc_request_id")
|
|
|
|
if not vnc_request_id:
|
|
logger.error("Missing vnc_request_id in stop request")
|
|
return
|
|
|
|
context = vnc_request_worker_map.get(vnc_request_id)
|
|
if not context:
|
|
logger.warning(f"No mapping found for stop request {vnc_request_id}")
|
|
return
|
|
|
|
worker_id = context["worker_id"]
|
|
socket_id = connected_workers.get(worker_id)
|
|
if socket_id:
|
|
emit("stop_vnc_stream_on_worker", {
|
|
"vnc_request_id": vnc_request_id,
|
|
}, to=socket_id)
|
|
logger.info(f"Sent stop_vnc_stream to worker {worker_id}")
|
|
else:
|
|
logger.warning(f"Worker {worker_id} not connected, skip sending stop signal")
|
|
|
|
# Clean up mapping
|
|
del vnc_request_worker_map[vnc_request_id]
|
|
logger.info(f"Cleaned up vnc_request_worker_map entry for {vnc_request_id}")
|
|
|
|
|