Files
blackbird/aibox/subsystems/executors.py
T
qing 177a9ac9a6 fix(aibox): close audit findings H1-H3/M1-M2 (command exec protection chain)
- H1: CommandLineTool.execute_streaming now purely delegates to injected
  executor; removes 150-line duplicate Popen loop that bypassed 1MB cap,
  cancel() and timeout protections
- H2: thread the blocking readline into a daemon reader thread; main loop
  polls queue with 0.1s ticks so timeout/1MB checks fire even for silent
  commands (verified: 3s timeout on 15s sleep, was 15s+ before)
- H3: stdin=DEVNULL on both Popen sites; stdin-waiting commands now hit
  EOF instead of deadlocking the tool thread forever
- M1: heredoc delimiter uniquification in write_file guidance (content
  line equal to AIBOX_EOF no longer truncates into shell injection)
- M2: 3600s total deadline on registry streaming forward loop

Verification: PoC regression 5/5 PASS (timeout/1MB/cancel/delegation),
full unit suite 264 passed. Audit report updated with Batch E results.
2026-08-17 17:28:41 +08:00

343 lines
13 KiB
Python

#!/usr/bin/env python3
"""
subsystems.executors - 命令执行器
抽离自 aibox.py,提供命令执行的抽象基类与 subprocess 实现。
- CommandExecutor : 抽象基类,定义 execute / execute_streaming / cancel
- SubprocessCommandExecutor : 基于 subprocess.Popen 的实现,支持超时/工作目录/
1MB 输出上限/线程安全取消
依赖:common (ConfigManager, constants, Logger) + 标准库
不反向依赖 aibox 任何符号。
"""
import os
import queue
import subprocess
import threading
import time
from abc import ABC, abstractmethod
from typing import Dict, Optional, Tuple
from common import ConfigManager, constants, Logger
class CommandExecutor(ABC):
"""命令执行器抽象基类 - 负责执行系统命令"""
def __init__(self, logger: Optional[Logger] = None):
self.logger = logger or Logger("CommandExecutor")
self.config = ConfigManager()
self.process = None
@abstractmethod
def execute(self, command: str, timeout: int = 30, **kwargs) -> str:
"""执行命令并返回结果(非流式)"""
pass
@abstractmethod
def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30, **kwargs):
"""流式执行命令,通过队列输出结果"""
pass
def is_enabled(self) -> bool:
"""检查命令执行是否启用"""
return constants.COMMANDS_ENABLED
def format_result(self, success: bool, message: str, data: Dict = None) -> str:
"""统一结果格式化"""
result = []
emoji = "✅" if success else "❌"
result.append(f"{emoji} {message}")
if data:
for key, value in data.items():
if isinstance(value, str) and "\n" in value:
result.append(f" {key}:")
for line in value.split("\n"):
result.append(f" {line}")
else:
result.append(f" {key}: {value}")
return "\n".join(result)
class SubprocessCommandExecutor(CommandExecutor):
"""基于subprocess的命令执行器 - 支持超时、工作目录、输出大小限制和取消"""
# 命令输出最大字节数(1MB),超出截断
MAX_OUTPUT_BYTES = 1 * 1024 * 1024
def __init__(self, logger=None):
super().__init__(logger)
self._current_process = None
self._process_lock = threading.Lock()
def cancel(self):
"""取消当前正在运行的命令(kill子进程)。线程安全。"""
with self._process_lock:
proc = self._current_process
if proc is not None:
try:
if proc.poll() is None:
proc.kill()
self.logger.info("已终止正在运行的子进程")
except Exception:
pass
def _read_with_limit(self, process, timeout: int, start_time: float, output_queue=None) -> Tuple[str, bool, bool]:
"""从进程读取输出,限制总大小和超时。返回 (output_str, timeout_occurred, truncated)。
H2 修复:阻塞 readline 在独立 reader 线程执行,主循环带超时轮询队列。
此前超时检查只在两次 readline 之间执行,无输出命令会让 readline
阻塞到进程自然退出,timeout 形同虚设(审计 PoC 实证)。
"""
output_chunks = []
total_bytes = 0
timeout_occurred = False
truncated = False
line_queue: queue.Queue = queue.Queue()
reader_done = threading.Event()
def _reader():
"""后台 reader 线程:阻塞 readline,行入队;EOF/异常结束"""
try:
for line in iter(process.stdout.readline, ""):
line_queue.put(line)
except Exception:
pass
finally:
reader_done.set()
reader = threading.Thread(target=_reader, daemon=True)
reader.start()
shell_dead_since = None # shell 进程已死但 reader 未收到 EOF 的起始时刻
while True:
# 检查超时(readline 阻塞不再影响本判断)
if timeout > 0 and time.time() - start_time > timeout:
process.terminate()
timeout_occurred = True
break
# 检查输出大小
if total_bytes >= self.MAX_OUTPUT_BYTES:
truncated = True
process.terminate()
break
try:
line = line_queue.get(timeout=0.1)
shell_dead_since = None
except queue.Empty:
if reader_done.is_set():
break
# shell 已被 kill/退出但 reader 仍阻塞(如 Windows shell=True
# 下 cancel 只杀死包装进程,孙进程仍持有管道写端导致无 EOF):
# 宽限 1 秒排空残留输出后放弃等待
if process.poll() is not None:
if shell_dead_since is None:
shell_dead_since = time.time()
elif time.time() - shell_dead_since > 1.0:
break
continue
line_bytes = len(line.encode('utf-8', errors='replace'))
if total_bytes + line_bytes > self.MAX_OUTPUT_BYTES and not truncated:
allowed = self.MAX_OUTPUT_BYTES - total_bytes
if allowed > 0:
output_chunks.append(line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='replace'))
total_bytes = self.MAX_OUTPUT_BYTES
truncated = True
process.terminate()
break
output_chunks.append(line)
total_bytes += line_bytes
if output_queue:
output_queue.put(("line", line.rstrip('\n') + "\n"))
return ''.join(output_chunks), timeout_occurred, truncated
def execute(self, command: str, timeout: int = 30, work_dir: str = None, env: Dict = None, **kwargs) -> str:
"""执行命令并返回完整结果(带输出大小限制)"""
if not self.is_enabled():
return "❌ 命令执行功能已禁用"
# 默认工作目录为workspace,避免在启动目录散落文件
if work_dir is None or work_dir == ".":
work_dir = constants.WORKSPACE_DIR
self.logger.info(f"执行命令: {command} (work_dir={work_dir})")
# 准备环境变量
exec_env = os.environ.copy()
if env:
exec_env.update(env)
process = None
try:
process = subprocess.Popen(
command,
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
stdin=subprocess.DEVNULL, # H3 修复:不挂 stdin,防等待输入命令永久挂起
text=True,
bufsize=1,
universal_newlines=True,
env=exec_env,
cwd=os.path.expanduser(work_dir)
)
# 保存进程引用以便cancel()可以kill它
with self._process_lock:
self._current_process = process
start_time = time.time()
output_str, timeout_occurred, truncated = self._read_with_limit(process, timeout, start_time)
# 确保进程终止
if process.poll() is None:
if timeout_occurred or truncated:
process.kill()
else:
process.terminate()
try:
process.wait(timeout=2)
except:
process.kill()
return_code = process.returncode if process.returncode is not None else -1
data = {
"输出": output_str.strip() if output_str.strip() else "(无输出)",
"返回码": return_code
}
if timeout_occurred:
return self.format_result(False, f"命令执行超时({timeout}秒)", data)
if truncated:
data["警告"] = f"输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断"
return self.format_result(False, "命令执行完成(输出已截断)", data)
return self.format_result(return_code == 0, "命令执行完成", data)
except subprocess.TimeoutExpired:
return f"❌ 命令执行超时({timeout}秒)"
except FileNotFoundError as e:
return f"❌ 命令执行错误: 命令未找到 - {str(e)}"
except PermissionError as e:
return f"❌ 命令执行错误: 权限不足 - {str(e)}"
except Exception as e:
return f"❌ 执行错误: {str(e)}"
finally:
if process and process.poll() is None:
process.kill()
# 清除当前进程引用
with self._process_lock:
self._current_process = None
def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30,
work_dir: str = None, env: Dict = None, **kwargs):
"""流式执行命令,支持超时控制、工作目录和输出大小限制"""
if not self.is_enabled():
output_queue.put(("line", "❌ 命令执行功能已禁用\n"))
output_queue.put(("complete", "❌ 命令执行功能已禁用"))
return
# 默认工作目录为workspace,避免在启动目录散落文件
if work_dir is None or work_dir == ".":
work_dir = constants.WORKSPACE_DIR
self.logger.info(f"流式执行命令: {command} (work_dir={work_dir})")
output_queue.put(("line", f"\n🔧 执行命令: {command}\n"))
output_queue.put(("line", f"⏱️ 超时设置: {timeout} 秒\n"))
output_queue.put(("line", f"📦 输出限制: {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB\n"))
output_queue.put(("line", f"📁 工作目录: {work_dir}\n"))
output_queue.put(("line", "-" * 50 + "\n"))
# 准备环境变量
exec_env = os.environ.copy()
if env:
exec_env.update(env)
process = None
try:
process = subprocess.Popen(
command,
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
stdin=subprocess.DEVNULL, # H3 修复:不挂 stdin,防等待输入命令永久挂起
text=True,
bufsize=1,
universal_newlines=True,
env=exec_env,
cwd=os.path.expanduser(work_dir)
)
self.process = process
# 保存进程引用以便cancel()可以kill它
with self._process_lock:
self._current_process = process
start_time = time.time()
# H2 修复:复用线程化读取(reader 线程 + 超时轮询),行经 output_queue 流式转发
output_str, timeout_occurred, truncated = self._read_with_limit(
process, timeout, start_time, output_queue=output_queue)
if timeout_occurred:
output_queue.put(("line", f"\n⏰ 命令执行超时({timeout}秒)\n"))
elif truncated:
output_queue.put(("line", f"\n📦 输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断\n"))
# 如果超时或截断,确保进程被终止
if (timeout_occurred or truncated) and process.poll() is None:
process.kill()
try:
process.wait(timeout=2)
except:
pass
return_code = process.returncode if process.returncode is not None else -1
output_queue.put(("line", "-" * 50 + "\n"))
if timeout_occurred:
output_queue.put(("line", f"⏰ 命令执行超时({timeout}秒),返回码: {return_code}\n"))
elif truncated:
output_queue.put(("line", f"📦 输出已截断,返回码: {return_code}\n"))
else:
output_queue.put(("line", f"✅ 命令执行完成,返回码: {return_code}\n"))
data = {
"输出": output_str.strip() if output_str.strip() else "(无输出)",
"返回码": return_code
}
if timeout_occurred:
result = self.format_result(False, f"命令执行超时({timeout}秒)", data)
elif truncated:
data["警告"] = f"输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断"
result = self.format_result(False, "命令执行完成(输出已截断)", data)
else:
result = self.format_result(return_code == 0, "命令执行完成", data)
output_queue.put(("complete", result))
except Exception as e:
output_queue.put(("line", f"❌ 执行错误: {str(e)}\n"))
output_queue.put(("complete", f"❌ 执行错误: {str(e)}"))
finally:
if process and process.poll() is None:
process.terminate()
try:
process.wait(timeout=2)
except:
process.kill()
self.process = None
with self._process_lock:
self._current_process = None