154 lines
5.1 KiB
Python
154 lines
5.1 KiB
Python
import os
|
|
import select
|
|
import subprocess
|
|
import threading
|
|
from flask import Blueprint, request, current_app, session
|
|
from flask_socketio import SocketIO, emit, join_room
|
|
from threading import Thread
|
|
import sys
|
|
from ..extensions import socketio
|
|
|
|
terminal_bp = Blueprint('terminal', __name__, url_prefix='/vsc-like')
|
|
|
|
# 存储终端进程
|
|
terminals = {}
|
|
|
|
def start_terminal_process(socket_id):
|
|
# 启动一个 bash 终端
|
|
process = subprocess.Popen(
|
|
['bash'],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True
|
|
)
|
|
|
|
# 监听进程的输出
|
|
while True:
|
|
output = process.stdout.readline()
|
|
if output:
|
|
emit('terminal_output', {'data': output}, room=socket_id)
|
|
if process.poll() is not None:
|
|
break
|
|
|
|
|
|
def start_terminal_process(socket_id):
|
|
"""
|
|
监听多个子进程的输出，并通过 WebSocket 发给前端
|
|
使用 select 进行非阻塞的 I/O 操作，并实现动态流量控制
|
|
"""
|
|
process = terminals[socket_id]
|
|
|
|
# 获取子进程的 stdout 和 stderr
|
|
stdout_fd = process.stdout
|
|
stderr_fd = process.stderr
|
|
|
|
# 将 stdout 和 stderr 添加到 select 监听的文件描述符列表中
|
|
fds = [stdout_fd, stderr_fd]
|
|
|
|
# 初始等待时间
|
|
timeout = 0.1
|
|
max_timeout = 2 # 设置最大超时时间为 2 秒
|
|
|
|
while True:
|
|
# 使用 select 监听多个文件描述符，等待数据准备好
|
|
readable, _, _ = select.select(fds, [], [], timeout)
|
|
|
|
if readable:
|
|
# 有数据可读时，重置超时时间为 0.1 秒
|
|
timeout = 0.1
|
|
|
|
for fd in readable:
|
|
# 读取进程的标准输出或错误输出
|
|
output = fd.readline()
|
|
if output:
|
|
emit('terminal_output', {'data': output}, room=socket_id)
|
|
else:
|
|
# 如果没有数据可读，则逐渐增加超时时间
|
|
timeout = min(timeout * 2, max_timeout)
|
|
|
|
# 检查进程是否仍在运行，如果不在运行则退出循环
|
|
if process.poll() is not None:
|
|
break
|
|
|
|
@socketio.on('connect', namespace='/vsc-like/terminal')
|
|
def connect_terminal():
|
|
"""
|
|
处理 WebSocket 连接
|
|
- 验证会话
|
|
- 使用指定的用户启动 bash 终端
|
|
- 启动监听子进程输出的线程
|
|
"""
|
|
socket_id = session.get('uuid')
|
|
user_id = session.get('user_id') # 从 session 获取用户 ID
|
|
path = session.get('path')
|
|
|
|
if not path or not user_id:
|
|
emit('error', {'message': 'Authentication failed, no path or user found for this session'})
|
|
return False # 拒绝连接
|
|
|
|
# 使用 sudo 启动终端进程，以指定的 user_id 用户身份执行
|
|
try:
|
|
process = subprocess.Popen(
|
|
['sudo', '-u', user_id, 'bash'], # 使用指定的 user_id 启动 bash
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True
|
|
)
|
|
terminals[socket_id] = process
|
|
|
|
emit('connected', {'message': 'Terminal connected'})
|
|
join_room(path, namespace='/vsc-like/terminal')
|
|
|
|
# 启动监听终端输出的线程
|
|
thread = Thread(target=start_terminal_process, args=(socket_id,))
|
|
thread.daemon = True
|
|
thread.start()
|
|
|
|
except Exception as e:
|
|
emit('error', {'message': f'Error starting terminal process: {str(e)}'})
|
|
return False
|
|
|
|
terminal_locks = {} # 存储终端的锁，用于线程同步
|
|
|
|
@socketio.on('send_input', namespace='/vsc-like/terminal')
|
|
def send_input(data):
|
|
socket_id = session.get('uuid')
|
|
input_data = data.get('input')
|
|
|
|
# 检查当前 socket_id 是否有对应的终端
|
|
if socket_id in terminals:
|
|
# 获取该终端的锁进行线程同步
|
|
lock = terminal_locks.get(socket_id)
|
|
if not lock:
|
|
lock = threading.Lock()
|
|
terminal_locks[socket_id] = lock
|
|
|
|
with lock: # 在锁的保护下操作终端
|
|
try:
|
|
terminals[socket_id].stdin.write(input_data + '\n')
|
|
terminals[socket_id].stdin.flush()
|
|
emit('terminal_output', {'data': f'Command sent: {input_data}'}, room=socket_id)
|
|
except Exception as e:
|
|
emit('error', {'message': f'Error sending input: {str(e)}'}, room=socket_id)
|
|
else:
|
|
emit('error', {'message': 'No terminal found for this session'}, room=socket_id)
|
|
|
|
|
|
@socketio.on('disconnect', namespace='/vsc-like/terminal')
|
|
def disconnect_terminal():
|
|
socket_id = session.get('uuid')
|
|
|
|
# 检查是否有终端进程需要终止
|
|
if socket_id in terminals:
|
|
try:
|
|
# 终止终端进程
|
|
terminals[socket_id].terminate()
|
|
del terminals[socket_id]
|
|
# 删除锁
|
|
if socket_id in terminal_locks:
|
|
del terminal_locks[socket_id]
|
|
emit('terminal_output', {'data': 'Terminal disconnected and terminated'}, room=socket_id)
|
|
except Exception as e:
|
|
emit('error', {'message': f'Error disconnecting terminal: {str(e)}'}, room=socket_id) |