159 lines
6.0 KiB
Python
159 lines
6.0 KiB
Python
import os
|
||
import select
|
||
import time
|
||
import uuid
|
||
import subprocess
|
||
import threading
|
||
from flask import Blueprint, request, current_app, session
|
||
from flask_socketio import Namespace, join_room, leave_room, emit, SocketIO
|
||
from threading import Thread
|
||
import sys
|
||
from ..extensions import socketio
|
||
|
||
terminal_bp = Blueprint('terminal', __name__, url_prefix='/vsc-like')
|
||
|
||
# 存储终端进程
|
||
command_dict = {}
|
||
terminal_busy = {}
|
||
terminals = {}
|
||
terminal_locks = {}
|
||
|
||
def check_process_running(command_name):
|
||
try:
|
||
# 使用 pgrep 查找与给定命令相关的进程
|
||
result = subprocess.run(
|
||
["pgrep", "-f", command_name],
|
||
stdout=subprocess.PIPE, stderr=subprocess.PIPE
|
||
)
|
||
return result.returncode == 0 # 如果返回值为 0，说明命令仍在运行
|
||
except Exception as e:
|
||
print(f"Error checking process: {e}")
|
||
return False
|
||
|
||
def monitor_process(socket_id, command_name, namespace):
|
||
while True:
|
||
is_running = check_process_running(command_name)
|
||
if is_running:
|
||
print(f"Process '{command_name}' is still running.")
|
||
else:
|
||
print(f"Process '{command_name}' is not running.")
|
||
break # 停止检查，如果进程不再运行
|
||
|
||
time.sleep(0.1) # 每 0.1 秒检查一次
|
||
command_dict[socket_id] = ''
|
||
terminal_busy[socket_id] = False
|
||
namespace.emit('terminal_output',{'data': '>'})
|
||
|
||
|
||
def start_terminal_process(socket_id, namespace):
|
||
"""
|
||
监听多个子进程的输出，并通过 WebSocket 发给前端
|
||
使用 select 进行非阻塞的 I/O 操作，并实现动态流量控制
|
||
"""
|
||
process = terminals[socket_id]
|
||
|
||
# 获取子进程的 stdout 和 stderr
|
||
stdout = process.stdout
|
||
stdin = process.stdin
|
||
# 将 stdout 和 stderr 添加到 select 监听的文件描述符列表中
|
||
|
||
# 初始等待时间
|
||
timeout = 0.1
|
||
max_timeout = 2 # 设置最大超时时间为 2 秒
|
||
|
||
while True:
|
||
while True:
|
||
output = stdout.readline()
|
||
if not output:
|
||
break
|
||
timeout = 0.1
|
||
namespace.emit('terminal_output', {'data': output+'\r'}, room=socket_id)
|
||
timeout = min(timeout * 2, max_timeout)
|
||
time.sleep(timeout)
|
||
if process.poll() is not None:
|
||
break
|
||
|
||
|
||
class VSCLikeNameSpace(Namespace):
|
||
def on_connect(self, data):
|
||
"""
|
||
处理 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', '-il'], # 使用指定的 user_id 启动 bash
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
text=True, # 确保输出是文本格式
|
||
)
|
||
terminals[socket_id] = process
|
||
terminal_busy[socket_id] = False
|
||
emit('connected', {'message': 'OK'})
|
||
join_room(socket_id, namespace='/vsc-like')
|
||
|
||
socketio.start_background_task(target=start_terminal_process, socket_id=socket_id, namespace=self)
|
||
|
||
except Exception as e:
|
||
emit('error', {'message': f'Error starting terminal process: {str(e)}'})
|
||
return False
|
||
|
||
def on_send_input(self, data):
|
||
socket_id = session.get('uuid')
|
||
input_data = data.get('input')
|
||
if socket_id not in command_dict:command_dict[socket_id]=''
|
||
# 检查当前 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)
|
||
command_dict[socket_id]+=input_data
|
||
if input_data=='\r':
|
||
terminals[socket_id].stdin.flush()
|
||
emit('terminal_output', {'data': f'\r\n'}, room=socket_id)
|
||
if terminal_busy[socket_id]==False:
|
||
terminal_busy[socket_id] = True
|
||
socketio.start_background_task(monitor_process, socket_id, command_dict[socket_id], namespace=self)
|
||
elif input_data=='\x03':
|
||
terminals[socket_id].stdin.flush()
|
||
emit('terminal_output', {'data': f'^C\x03\r\n'}, room=socket_id)
|
||
else:
|
||
emit('terminal_output', {'data': 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)
|
||
|
||
def on_disconnect():
|
||
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)
|