103 lines
3.8 KiB
Python
103 lines
3.8 KiB
Python
# websocket_server/events/terminal.py
|
|
|
|
from flask_socketio import emit
|
|
import json
|
|
import logging
|
|
import requests
|
|
|
|
from websocket_server.config import get_redis_client
|
|
from websocket_server.events import base
|
|
from websocket_server.shared_state import connected_workers, connected_sids_lock, connected_sids
|
|
|
|
logger = logging.getLogger("websocket_server")
|
|
socketio = base.socketio
|
|
|
|
api_server_url = "http://127.0.0.1:5000/api"
|
|
def register_socketio_handlers(socketio):
|
|
|
|
@socketio.on("start_terminal_session")
|
|
def handle_start_terminal_session(data):
|
|
container_id = data["container_id"]
|
|
request_id = data["request_id"]
|
|
user_sid = request.sid
|
|
|
|
logger.info(f"[start_terminal_session] Request ID: {request_id}, Container ID: {container_id}, User SID: {user_sid}")
|
|
|
|
try:
|
|
# Lookup container details
|
|
response = requests.get(f"{api_server_url}/workloads/containers/{container_id}")
|
|
response.raise_for_status()
|
|
info = response.json()
|
|
worker_id = info["workload_host_id"]
|
|
|
|
context = {
|
|
"user_sid": user_sid,
|
|
"worker_id": worker_id,
|
|
"container_id": container_id
|
|
}
|
|
|
|
redis_client = get_redis_client()
|
|
redis_key = f"terminal_context:{request_id}"
|
|
redis_client.setex(redis_key, 600, json.dumps(context))
|
|
logger.info(f"[start_terminal_session] Stored context in Redis under key {redis_key}")
|
|
|
|
if worker_id not in connected_workers:
|
|
logger.error(f"[start_terminal_session] Worker ID {worker_id} not connected.")
|
|
return
|
|
|
|
socketio.emit("start_terminal", {
|
|
"container_name": container_id,
|
|
"request_id": request_id
|
|
}, to=connected_workers[worker_id])
|
|
logger.info(f"[start_terminal_session] Emitted 'start_terminal' to worker {worker_id}")
|
|
|
|
except Exception as e:
|
|
logger.exception(f"[start_terminal_session] Exception occurred: {e}")
|
|
|
|
@socketio.on("stop_terminal_session")
|
|
def handle_stop_terminal_session(data):
|
|
request_id = data["request_id"]
|
|
redis_key = f"terminal_context:{request_id}"
|
|
|
|
try:
|
|
redis_client = get_redis_client()
|
|
context_raw = redis_client.get(redis_key)
|
|
context = json.loads(context_raw or '{}')
|
|
|
|
if not context:
|
|
logger.warning(f"[stop_terminal_session] No context found for request ID {request_id}")
|
|
return
|
|
|
|
worker_id = context["worker_id"]
|
|
socketio.emit("stop_terminal", {"request_id": request_id}, to=connected_workers[worker_id])
|
|
logger.info(f"[stop_terminal_session] Emitted 'stop_terminal' to worker {worker_id}")
|
|
|
|
redis_client.delete(redis_key)
|
|
logger.info(f"[stop_terminal_session] Deleted Redis key {redis_key}")
|
|
|
|
except Exception as e:
|
|
logger.exception(f"[stop_terminal_session] Exception occurred: {e}")
|
|
|
|
|
|
@socketio.on("terminal_input")
|
|
def handle_terminal_input(data):
|
|
request_id = data["request_id"]
|
|
context = json.loads(get_redis_client().get(f"terminal_context:{request_id}") or '{}')
|
|
if context:
|
|
socketio.emit("terminal_input", {
|
|
"request_id": request_id,
|
|
"data": data["data"]
|
|
}, to=connected_workers[context["worker_id"]])
|
|
|
|
|
|
@socketio.on("terminal_output")
|
|
def handle_terminal_output(data):
|
|
request_id = data["request_id"]
|
|
context = json.loads(get_redis_client().get(f"terminal_context:{request_id}") or '{}')
|
|
if context:
|
|
socketio.emit("terminal_output", {
|
|
"request_id": request_id,
|
|
"data": data["data"]
|
|
}, to=context["user_sid"])
|
|
|