mirror of
https://github.com/yyhuni/xingrin.git
synced 2026-01-31 11:46:16 +08:00
807 lines
30 KiB
Python
807 lines
30 KiB
Python
"""
|
||
命令执行器
|
||
|
||
统一管理所有命令执行方式:
|
||
- execute_and_wait(): 等待式执行,适合输出到文件的工具
|
||
- execute_stream(): 流式执行,适合实时处理输出的工具
|
||
|
||
性能监控:
|
||
- 自动记录命令执行耗时、内存使用
|
||
- 输出到 performance logger
|
||
"""
|
||
|
||
import logging
|
||
import os
|
||
import re
|
||
import signal
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
from collections import deque
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Dict, Any, Optional, Generator
|
||
|
||
from django.conf import settings
|
||
|
||
try:
|
||
# 可选依赖:用于根据 CPU / 内存负载做动态并发控制
|
||
import psutil
|
||
except ImportError: # 运行环境缺少 psutil 时降级为无动态负载控制
|
||
psutil = None
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 延迟导入,避免循环依赖
|
||
def _get_command_tracker(tool_name: str, command: str):
|
||
"""获取命令性能追踪器(延迟导入)"""
|
||
from apps.scan.utils.performance import CommandPerformanceTracker
|
||
return CommandPerformanceTracker(tool_name, command)
|
||
|
||
# 常量定义
|
||
GRACEFUL_SHUTDOWN_TIMEOUT = 5 # 进程优雅退出的超时时间(秒)
|
||
MAX_LOG_TAIL_LINES = 1000 # 日志文件读取的最大行数
|
||
|
||
# 命令日志配置(从环境变量读取)
|
||
# ENABLE_COMMAND_LOGGING=true: 输出所有内容(命令输出+错误)到log_file_path
|
||
# ENABLE_COMMAND_LOGGING=false: 只输出错误到log_file_path
|
||
ENABLE_COMMAND_LOGGING = getattr(settings, 'ENABLE_COMMAND_LOGGING', True)
|
||
|
||
# 动态并发控制阈值(可在 Django settings 中覆盖)
|
||
SCAN_CPU_HIGH = getattr(settings, 'SCAN_CPU_HIGH', 90.0) # CPU 高水位(百分比)
|
||
SCAN_MEM_HIGH = getattr(settings, 'SCAN_MEM_HIGH', 80.0) # 内存高水位(百分比)
|
||
SCAN_LOAD_CHECK_INTERVAL = getattr(settings, 'SCAN_LOAD_CHECK_INTERVAL', 180) # 负载检查间隔(秒)
|
||
SCAN_COMMAND_STARTUP_DELAY = getattr(settings, 'SCAN_COMMAND_STARTUP_DELAY', 5) # 命令启动前等待(秒)
|
||
|
||
_ACTIVE_COMMANDS = 0
|
||
_ACTIVE_COMMANDS_LOCK = threading.Lock()
|
||
|
||
|
||
def _wait_for_system_load() -> None:
|
||
"""根据当前机器 CPU/内存负载,决定是否暂缓启动新的外部命令。"""
|
||
|
||
# 1. 先强制等待,让之前启动的命令有时间消耗资源,避免并发启动导致延迟OOM
|
||
if SCAN_COMMAND_STARTUP_DELAY > 0:
|
||
time.sleep(SCAN_COMMAND_STARTUP_DELAY)
|
||
|
||
# 2. 再检查系统负载
|
||
if psutil is None:
|
||
raise ImportError("psutil 未安装,无法进行负载感知控制")
|
||
|
||
while True:
|
||
cpu = psutil.cpu_percent(interval=0.5)
|
||
mem = psutil.virtual_memory().percent
|
||
|
||
if cpu < SCAN_CPU_HIGH and mem < SCAN_MEM_HIGH:
|
||
return
|
||
|
||
logger.info(
|
||
"系统负载较高,任务将排队执行,防止oom: cpu=%.1f%% (阈值 %.1f%%), mem=%.1f%% (阈值 %.1f%%)",
|
||
cpu,
|
||
SCAN_CPU_HIGH,
|
||
mem,
|
||
SCAN_MEM_HIGH,
|
||
)
|
||
time.sleep(SCAN_LOAD_CHECK_INTERVAL)
|
||
|
||
|
||
class CommandExecutor:
|
||
"""
|
||
统一的命令执行器
|
||
|
||
提供两种执行模式:
|
||
1. execute_and_wait() - 等待式执行(适合文件输出)
|
||
2. execute_stream() - 流式执行(适合实时处理)
|
||
"""
|
||
|
||
def _write_command_start_header(self, log_file: Path, tool_name: str, command: str, timeout: Optional[int] = None):
|
||
"""
|
||
在命令开始时写入头部信息
|
||
"""
|
||
if not ENABLE_COMMAND_LOGGING:
|
||
return
|
||
|
||
try:
|
||
with open(log_file, 'w', encoding='utf-8') as f:
|
||
f.write(f"$ {command}\n")
|
||
f.write(f"{'='*60}\n")
|
||
f.write(f"# 工具: {tool_name}\n")
|
||
f.write(f"# 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
|
||
if timeout is not None:
|
||
f.write(f"# 超时限制: {timeout}秒\n")
|
||
f.write(f"# 状态: 执行中...\n")
|
||
f.write(f"{'='*60}\n\n")
|
||
except Exception as e:
|
||
logger.error(f"写入命令开始信息失败: {e}")
|
||
|
||
def _write_command_end_footer(self, log_file: Path, tool_name: str, duration: float, returncode: int, success: bool):
|
||
"""
|
||
在命令结束时追加尾部信息
|
||
"""
|
||
if not ENABLE_COMMAND_LOGGING:
|
||
return
|
||
|
||
try:
|
||
with open(log_file, 'a', encoding='utf-8') as f:
|
||
f.write(f"\n{'='*60}\n")
|
||
f.write(f"# 结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
|
||
f.write(f"# 执行耗时: {duration:.2f}秒\n")
|
||
f.write(f"# 退出码: {returncode}\n")
|
||
f.write(f"# 状态: {'✓ 成功' if success else '✗ 失败'}\n")
|
||
f.write(f"{'='*60}\n")
|
||
|
||
logger.info(f"📝 {tool_name} 日志: {log_file} (耗时: {duration:.2f}秒)")
|
||
except Exception as e:
|
||
logger.error(f"写入命令结束信息失败: {e}")
|
||
|
||
def _clean_output_line(self, line: str, suffix_char: Optional[str] = None) -> Optional[str]:
|
||
"""
|
||
统一的输出行清理处理
|
||
|
||
处理顺序:
|
||
1. 去除首尾空白
|
||
2. 跳过空行
|
||
3. 处理字面转义字符串(如 \\x0d\\x0a)
|
||
4. 移除 ANSI 转义序列
|
||
5. 清理控制字符
|
||
6. 移除指定后缀字符
|
||
|
||
Args:
|
||
line: 原始输出行
|
||
suffix_char: 要移除的末尾字符
|
||
|
||
Returns:
|
||
清理后的行内容,如果是空行则返回 None
|
||
"""
|
||
# 1. 去除行首尾的空白字符
|
||
line = line.strip()
|
||
|
||
# 2. 跳过空行
|
||
if not line:
|
||
return None
|
||
|
||
# 3. 处理字面转义字符串(罕见但可能存在)
|
||
# 处理常见的字面转义序列
|
||
escape_mappings = {
|
||
'\\x0d\\x0a': '\n', # Windows 换行符字面量
|
||
'\\x0a': '\n', # Unix 换行符字面量
|
||
'\\x0d': '\r', # 回车符字面量
|
||
'\\r\\n': '\n', # 常见的转义表示
|
||
'\\n': '\n', # 换行符转义
|
||
'\\r': '\r', # 回车符转义
|
||
'\\t': '\t', # 制表符转义
|
||
}
|
||
|
||
for literal, actual in escape_mappings.items():
|
||
if literal in line:
|
||
line = line.replace(literal, actual)
|
||
|
||
# 4. 移除 ANSI 转义序列(颜色、格式等控制字符)
|
||
ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
|
||
line = ansi_escape.sub('', line)
|
||
|
||
# 5. 清理控制字符
|
||
line = line.replace('\x00', '') # 移除 NUL 字符
|
||
line = line.replace('\r', '') # 移除回车符
|
||
line = line.replace('\b', '') # 移除退格符
|
||
line = line.replace('\f', '') # 移除换页符
|
||
line = line.replace('\v', '') # 移除垂直制表符
|
||
|
||
# 6. 再次去除空白(清理后可能产生新的空白)
|
||
line = line.strip()
|
||
if not line:
|
||
return None
|
||
|
||
# 7. 如果指定了后缀字符,移除末尾的后缀字符
|
||
if suffix_char and line.endswith(suffix_char):
|
||
line = line[:-1].strip()
|
||
if not line:
|
||
return None
|
||
|
||
return line
|
||
|
||
def _kill_process_tree(self, process: subprocess.Popen) -> None:
|
||
"""
|
||
强制终止进程树
|
||
|
||
当使用 shell=True 时,process.pid 是 shell 的 PID。
|
||
如果不杀掉整个进程组,shell 的子进程(实际工具)会变成孤儿进程继续运行。
|
||
"""
|
||
if process.poll() is not None:
|
||
return
|
||
|
||
try:
|
||
# 尝试杀掉进程组(需要进程启动时设置 start_new_session=True)
|
||
os.killpg(os.getpgid(process.pid), signal.SIGKILL)
|
||
logger.debug(f"已终止进程组: PGID={process.pid}")
|
||
except ProcessLookupError:
|
||
pass # 进程已不存在
|
||
except Exception as e:
|
||
logger.warning(f"终止进程组失败 ({e}),尝试普通 kill")
|
||
try:
|
||
process.kill()
|
||
except Exception:
|
||
pass
|
||
|
||
def execute_and_wait(
|
||
self,
|
||
tool_name: str,
|
||
command: str,
|
||
timeout: int,
|
||
log_file: Optional[str] = None
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
等待式执行:启动命令并等待完成
|
||
|
||
适用场景:工具输出到文件(如 subfinder -o output.txt)
|
||
|
||
Args:
|
||
tool_name: 工具名称(用于日志)
|
||
command: 完整的扫描命令(包含输出文件参数)
|
||
timeout: 超时时间(秒)
|
||
log_file: 日志文件路径(可选,None 表示丢弃 stderr)
|
||
|
||
Returns:
|
||
dict: {
|
||
'success': bool, # 命令是否成功执行(returncode == 0)
|
||
'returncode': int, # 命令退出码
|
||
'log_file': str | None # 日志文件路径
|
||
}
|
||
|
||
Raises:
|
||
ValueError: 参数验证失败
|
||
RuntimeError: 执行失败或超时
|
||
"""
|
||
global _ACTIVE_COMMANDS
|
||
|
||
# 验证参数
|
||
if not tool_name:
|
||
raise ValueError("工具名称不能为空")
|
||
if not command:
|
||
raise ValueError("扫描命令不能为空")
|
||
if timeout <= 0:
|
||
raise ValueError(f"超时时间必须大于0: {timeout}")
|
||
|
||
|
||
logger.info("开始运行扫描工具: %s", tool_name)
|
||
|
||
# 准备日志文件
|
||
log_file_path = Path(log_file) if log_file else None
|
||
|
||
# 记录开始时间(用于计算执行时间)
|
||
start_time = datetime.now()
|
||
|
||
# 初始化性能追踪器
|
||
perf_tracker = _get_command_tracker(tool_name, command)
|
||
perf_tracker.start()
|
||
|
||
process = None
|
||
log_file_handle = None
|
||
acquired_slot = False # 标记是否已增加全局活动命令计数
|
||
|
||
try:
|
||
# 在启动新的外部命令之前,先根据 CPU/内存负载判断是否需要等待
|
||
_wait_for_system_load()
|
||
|
||
acquired_slot = True
|
||
if _ACTIVE_COMMANDS_LOCK:
|
||
with _ACTIVE_COMMANDS_LOCK:
|
||
_ACTIVE_COMMANDS += 1
|
||
current_active = _ACTIVE_COMMANDS
|
||
else:
|
||
current_active = 0
|
||
logger.info(
|
||
"登记活动命令计数: tool=%s, active=%d",
|
||
tool_name,
|
||
current_active,
|
||
)
|
||
|
||
logger.debug("执行命令: %s", command)
|
||
if log_file_path:
|
||
logger.debug("日志文件: %s", log_file_path)
|
||
else:
|
||
logger.debug("日志输出: 丢弃")
|
||
|
||
# 准备输出流
|
||
stdout_target = subprocess.DEVNULL
|
||
stderr_target = subprocess.DEVNULL
|
||
|
||
if log_file_path:
|
||
# 先写入命令开始信息
|
||
if ENABLE_COMMAND_LOGGING:
|
||
self._write_command_start_header(log_file_path, tool_name, command, timeout)
|
||
|
||
# 以追加模式打开日志文件
|
||
log_file_handle = open(log_file_path, 'a', encoding='utf-8', buffering=1)
|
||
if ENABLE_COMMAND_LOGGING:
|
||
stdout_target = log_file_handle
|
||
stderr_target = subprocess.STDOUT
|
||
else:
|
||
stderr_target = log_file_handle
|
||
|
||
# 启动进程
|
||
# 使用 start_new_session=True 创建新会话,使子进程成为新进程组的首领
|
||
# 这样我们可以通过 killpg 杀掉整个进程树
|
||
process = subprocess.Popen(
|
||
command,
|
||
stdin=subprocess.DEVNULL,
|
||
shell=True,
|
||
stdout=stdout_target,
|
||
stderr=stderr_target,
|
||
text=True,
|
||
start_new_session=True
|
||
)
|
||
|
||
# 设置进程 PID 用于性能追踪
|
||
perf_tracker.set_pid(process.pid)
|
||
|
||
# 等待完成
|
||
process.communicate(timeout=timeout)
|
||
|
||
# 检查执行结果
|
||
returncode = process.returncode
|
||
success = (returncode == 0)
|
||
|
||
# 计算执行时间
|
||
duration = (datetime.now() - start_time).total_seconds()
|
||
|
||
# 追加命令结束信息(如果开启且有日志文件)
|
||
if log_file_path and ENABLE_COMMAND_LOGGING:
|
||
self._write_command_end_footer(log_file_path, tool_name, duration, returncode, success)
|
||
command_log_file = str(log_file_path) if log_file_path else None
|
||
|
||
if not success:
|
||
# 命令执行失败,尝试读取错误日志
|
||
error_output = ""
|
||
if log_file_path:
|
||
error_output = self._read_log_tail(log_file_path, max_lines=MAX_LOG_TAIL_LINES)
|
||
logger.warning(
|
||
"扫描工具 %s 返回非零状态码: %d (执行时间: %.2f秒)",
|
||
tool_name, returncode, duration
|
||
)
|
||
if error_output:
|
||
for line in error_output.strip().split('\n'):
|
||
if line.strip():
|
||
logger.warning("%s", line)
|
||
else:
|
||
logger.info("✓ 扫描工具 %s 执行完成 (执行时间: %.2f秒)", tool_name, duration)
|
||
|
||
# 记录性能日志
|
||
perf_tracker.finish(success=success, duration=duration, timeout=timeout)
|
||
|
||
return {
|
||
'success': success,
|
||
'returncode': returncode,
|
||
'log_file': str(log_file_path) if log_file_path else None,
|
||
'command_log_file': command_log_file,
|
||
'duration': duration
|
||
}
|
||
|
||
except subprocess.TimeoutExpired as e:
|
||
# 计算超时时的执行时间
|
||
duration = (datetime.now() - start_time).total_seconds()
|
||
|
||
# 追加超时结束信息
|
||
if log_file_path and ENABLE_COMMAND_LOGGING:
|
||
self._write_command_end_footer(log_file_path, tool_name, duration, -1, False)
|
||
|
||
# 记录性能日志(超时)
|
||
perf_tracker.finish(success=False, duration=duration, timeout=timeout, is_timeout=True)
|
||
|
||
error_msg = f"扫描工具 {tool_name} 执行超时({timeout}秒,实际执行: {duration:.2f}秒)"
|
||
logger.error(error_msg)
|
||
if log_file_path and log_file_path.exists():
|
||
logger.debug("超时日志已保存: %s", log_file_path)
|
||
raise RuntimeError(error_msg) from e
|
||
|
||
except subprocess.SubprocessError as e:
|
||
error_msg = f"扫描工具 {tool_name} 执行失败: {e}"
|
||
logger.error(error_msg)
|
||
raise RuntimeError(error_msg) from e
|
||
|
||
except Exception as e:
|
||
# 捕获所有异常(包括 Prefect 取消引发的 CancelledError 等)
|
||
# 确保在 finally 块中清理进程
|
||
error_msg = f"扫描工具 {tool_name} 执行异常(可能是被中断): {e}"
|
||
logger.error(error_msg, exc_info=True)
|
||
raise
|
||
|
||
finally:
|
||
# 关键修复:确保进程树被清理
|
||
if process:
|
||
self._kill_process_tree(process)
|
||
|
||
# 关闭文件句柄
|
||
if log_file_handle:
|
||
try:
|
||
log_file_handle.close()
|
||
except Exception:
|
||
pass
|
||
|
||
if acquired_slot:
|
||
if _ACTIVE_COMMANDS_LOCK:
|
||
with _ACTIVE_COMMANDS_LOCK:
|
||
if _ACTIVE_COMMANDS > 0:
|
||
_ACTIVE_COMMANDS -= 1
|
||
current_active = _ACTIVE_COMMANDS
|
||
else:
|
||
current_active = 0
|
||
logger.info(
|
||
"释放活动命令计数: tool=%s, active=%d",
|
||
tool_name,
|
||
current_active,
|
||
)
|
||
|
||
def execute_stream(
|
||
self,
|
||
cmd: str,
|
||
tool_name: str,
|
||
cwd: Optional[str] = None,
|
||
shell: bool = False,
|
||
encoding: str = 'utf-8',
|
||
suffix_char: Optional[str] = None,
|
||
timeout: Optional[int] = None,
|
||
log_file: Optional[str] = None
|
||
) -> Generator[str, None, None]:
|
||
"""
|
||
流式执行:逐行返回输出
|
||
|
||
适用场景:工具流式输出 JSON(如 naabu -json)
|
||
|
||
Args:
|
||
cmd: 要执行的命令
|
||
tool_name: 工具名称(用于日志记录)
|
||
cwd: 工作目录
|
||
shell: 是否使用 shell 执行
|
||
encoding: 编码格式
|
||
suffix_char: 末尾后缀字符(用于移除)
|
||
timeout: 命令执行超时时间(秒),None 表示不设置超时
|
||
log_file: 日志文件路径(可选)
|
||
|
||
Yields:
|
||
str: 每行输出的内容(已处理:去空白、去ANSI、去后缀)
|
||
|
||
Raises:
|
||
subprocess.TimeoutExpired: 命令执行超时
|
||
"""
|
||
|
||
global _ACTIVE_COMMANDS
|
||
|
||
# 记录开始时间(用于命令日志)
|
||
start_time = datetime.now()
|
||
acquired_slot = False
|
||
|
||
# 初始化性能追踪器
|
||
perf_tracker = _get_command_tracker(tool_name, cmd)
|
||
perf_tracker.start()
|
||
|
||
# 准备日志文件路径
|
||
log_file_path = Path(log_file) if log_file else None
|
||
if log_file_path:
|
||
logger.debug(f"日志文件: {log_file_path}")
|
||
else:
|
||
logger.debug("日志输出: 丢弃")
|
||
|
||
# 根据是否使用shell来格式化命令
|
||
command = cmd if shell else cmd.split()
|
||
|
||
# 日志文件句柄
|
||
log_file_handle = None
|
||
|
||
# 启动子进程,根据日志策略决定输出方向
|
||
if log_file_path:
|
||
# 先写入命令开始信息
|
||
if ENABLE_COMMAND_LOGGING:
|
||
self._write_command_start_header(log_file_path, tool_name, cmd, timeout)
|
||
|
||
# 以追加模式打开日志文件(开始信息已写入)
|
||
log_file_handle = open(log_file_path, 'a', encoding='utf-8', buffering=1)
|
||
|
||
stdout_target = subprocess.PIPE
|
||
stderr_target = log_file_handle
|
||
if ENABLE_COMMAND_LOGGING:
|
||
stderr_target = subprocess.STDOUT
|
||
|
||
if not acquired_slot:
|
||
# 日志模式下,在真正启动进程前做一次负载检查,并登记活动命令计数
|
||
_wait_for_system_load()
|
||
acquired_slot = True
|
||
if _ACTIVE_COMMANDS_LOCK:
|
||
with _ACTIVE_COMMANDS_LOCK:
|
||
_ACTIVE_COMMANDS += 1
|
||
current_active = _ACTIVE_COMMANDS
|
||
else:
|
||
current_active = 0
|
||
logger.info(
|
||
"登记活动命令计数: tool=%s, active=%d",
|
||
tool_name,
|
||
current_active,
|
||
)
|
||
|
||
process = subprocess.Popen(
|
||
command,
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=stdout_target,
|
||
stderr=stderr_target,
|
||
cwd=cwd,
|
||
universal_newlines=True,
|
||
encoding=encoding,
|
||
shell=shell,
|
||
start_new_session=True # 关键:创建新进程组
|
||
)
|
||
else:
|
||
# 无日志文件:正常流式输出
|
||
if not acquired_slot:
|
||
# 非日志模式,同样在启动进程前做一次负载检查,并登记活动命令计数
|
||
_wait_for_system_load()
|
||
acquired_slot = True
|
||
if _ACTIVE_COMMANDS_LOCK:
|
||
with _ACTIVE_COMMANDS_LOCK:
|
||
_ACTIVE_COMMANDS += 1
|
||
current_active = _ACTIVE_COMMANDS
|
||
else:
|
||
current_active = 0
|
||
logger.info(
|
||
"登记活动命令计数: tool=%s, active=%d",
|
||
tool_name,
|
||
current_active,
|
||
)
|
||
|
||
process = subprocess.Popen(
|
||
command,
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
cwd=cwd,
|
||
universal_newlines=True,
|
||
encoding=encoding,
|
||
shell=shell,
|
||
start_new_session=True # 关键:创建新进程组
|
||
)
|
||
|
||
# 设置进程 PID 用于性能追踪
|
||
perf_tracker.set_pid(process.pid)
|
||
|
||
# 超时控制:使用 Timer 在指定时间后终止进程
|
||
timed_out_event = threading.Event()
|
||
|
||
def _kill_when_timeout():
|
||
timed_out_event.set()
|
||
if process.poll() is None: # 进程还在运行
|
||
logger.warning(f"命令执行超时({timeout}秒),正在终止进程: {cmd}")
|
||
self._kill_process_tree(process) # 使用新的终止方法
|
||
|
||
timer = None
|
||
if timeout is not None:
|
||
timer = threading.Timer(timeout, _kill_when_timeout)
|
||
timer.start()
|
||
|
||
try:
|
||
# 逐行读取进程输出
|
||
stdout = process.stdout
|
||
assert stdout is not None, "stdout should not be None when stdout=PIPE"
|
||
|
||
for line in iter(lambda: stdout.readline(), ''):
|
||
if not line:
|
||
break
|
||
|
||
# 统一字符处理
|
||
cleaned_line = self._clean_output_line(line, suffix_char)
|
||
if cleaned_line is None:
|
||
continue # 跳过空行
|
||
line = cleaned_line
|
||
|
||
# 如果开启命令日志且有日志文件,同时写入日志文件
|
||
if log_file_handle and ENABLE_COMMAND_LOGGING:
|
||
log_file_handle.write(line + '\n')
|
||
log_file_handle.flush()
|
||
|
||
# 直接返回行内容,由调用者负责解析
|
||
yield line
|
||
|
||
finally:
|
||
# 1. 停止定时器(如果还没触发)
|
||
if timer:
|
||
timer.cancel()
|
||
timer.join(timeout=0.1) # 等待 timer 线程完全结束,避免悬挂
|
||
|
||
# 2. 清理进程资源
|
||
exit_code = None
|
||
|
||
if timed_out_event.is_set():
|
||
# 超时情况:定时器已经处理了进程终止,只需获取退出码
|
||
logger.debug("进程已被超时定时器终止,等待进程结束")
|
||
try:
|
||
exit_code = process.wait(timeout=1.0) # 等待进程完全退出
|
||
except subprocess.TimeoutExpired:
|
||
logger.warning("进程在超时后仍未退出,强制终止")
|
||
self._kill_process_tree(process)
|
||
exit_code = -1
|
||
else:
|
||
# 正常结束:等待进程自然结束
|
||
# 如果是被外部中断(如 CancelledError),poll() 应为 None,需要 kill
|
||
if process.poll() is None:
|
||
logger.info(f"流式执行被中断,清理进程: {tool_name}")
|
||
self._kill_process_tree(process)
|
||
|
||
try:
|
||
exit_code = process.wait(timeout=GRACEFUL_SHUTDOWN_TIMEOUT)
|
||
except subprocess.TimeoutExpired:
|
||
logger.warning(
|
||
"程序未能在%d秒内自然结束,强制终止: %s",
|
||
GRACEFUL_SHUTDOWN_TIMEOUT, cmd
|
||
)
|
||
self._kill_process_tree(process)
|
||
exit_code = -2
|
||
|
||
# 3. 关闭进程流
|
||
if process.stdout:
|
||
process.stdout.close()
|
||
if process.stderr:
|
||
process.stderr.close()
|
||
|
||
# 4. 关闭日志文件句柄
|
||
if log_file_handle:
|
||
log_file_handle.close()
|
||
|
||
# 5. 追加命令结束信息(如果开启且有日志文件)
|
||
duration = (datetime.now() - start_time).total_seconds()
|
||
success = not timed_out_event.is_set() and (exit_code == 0 if exit_code is not None else True)
|
||
|
||
if log_file_path and ENABLE_COMMAND_LOGGING:
|
||
# 追加结束信息到日志文件末尾
|
||
self._write_command_end_footer(log_file_path, tool_name, duration, exit_code or 0, success)
|
||
|
||
# 6. 记录性能日志
|
||
perf_tracker.finish(success=success, duration=duration, timeout=timeout, is_timeout=timed_out_event.is_set())
|
||
|
||
if acquired_slot:
|
||
if _ACTIVE_COMMANDS_LOCK:
|
||
with _ACTIVE_COMMANDS_LOCK:
|
||
if _ACTIVE_COMMANDS > 0:
|
||
_ACTIVE_COMMANDS -= 1
|
||
current_active = _ACTIVE_COMMANDS
|
||
else:
|
||
current_active = 0
|
||
logger.info(
|
||
"释放活动命令计数: tool=%s, active=%d",
|
||
tool_name,
|
||
current_active,
|
||
)
|
||
|
||
def _read_log_tail(self, log_file: Path, max_lines: int = MAX_LOG_TAIL_LINES) -> str:
|
||
"""
|
||
读取日志文件的末尾部分(常量内存实现)
|
||
|
||
使用 seek 从文件末尾往前读取,避免将整个文件加载到内存。
|
||
|
||
Args:
|
||
log_file: 日志文件路径
|
||
max_lines: 最大读取行数
|
||
|
||
Returns:
|
||
日志内容(字符串),读取失败返回错误提示
|
||
"""
|
||
if not log_file.exists():
|
||
logger.debug("日志文件不存在: %s", log_file)
|
||
return ""
|
||
|
||
file_size = log_file.stat().st_size
|
||
if file_size == 0:
|
||
logger.debug("日志文件为空: %s", log_file)
|
||
return ""
|
||
|
||
# 每次读取的块大小(8KB,足够容纳大多数日志行)
|
||
chunk_size = 8192
|
||
|
||
def decode_line(line_bytes: bytes) -> str:
|
||
"""解码单行:优先 UTF-8,失败则降级 latin-1"""
|
||
try:
|
||
return line_bytes.decode('utf-8')
|
||
except UnicodeDecodeError:
|
||
return line_bytes.decode('latin-1', errors='replace')
|
||
|
||
try:
|
||
with open(log_file, 'rb') as f:
|
||
lines_found: deque[bytes] = deque()
|
||
remaining = b''
|
||
position = file_size
|
||
|
||
while position > 0 and len(lines_found) < max_lines:
|
||
read_size = min(chunk_size, position)
|
||
position -= read_size
|
||
|
||
f.seek(position)
|
||
chunk = f.read(read_size) + remaining
|
||
parts = chunk.split(b'\n')
|
||
|
||
# 最前面的部分可能不完整,留到下次处理
|
||
remaining = parts[0]
|
||
|
||
# 其余部分是完整的行(从后往前收集,用 appendleft 保持顺序)
|
||
for part in reversed(parts[1:]):
|
||
if len(lines_found) >= max_lines:
|
||
break
|
||
lines_found.appendleft(part)
|
||
|
||
# 处理文件开头的行
|
||
if remaining and len(lines_found) < max_lines:
|
||
lines_found.appendleft(remaining)
|
||
|
||
return '\n'.join(decode_line(line) for line in lines_found)
|
||
|
||
except PermissionError as e:
|
||
logger.warning("日志文件权限不足 (%s): %s", log_file, e)
|
||
return "(无法读取日志文件: 权限不足)"
|
||
except IOError as e:
|
||
logger.warning("日志文件读取IO错误 (%s): %s", log_file, e)
|
||
return f"(无法读取日志文件: IO错误 - {e})"
|
||
except Exception as e:
|
||
logger.warning("读取日志文件失败 (%s): %s", log_file, e, exc_info=True)
|
||
return f"(无法读取日志文件: {type(e).__name__} - {e})"
|
||
|
||
|
||
# 单例实例
|
||
_executor = CommandExecutor()
|
||
|
||
|
||
# 快捷函数
|
||
def execute_and_wait(
|
||
tool_name: str,
|
||
command: str,
|
||
timeout: int,
|
||
log_file: Optional[str] = None
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
等待式执行命令(快捷函数)
|
||
|
||
适用场景:工具输出到文件(如 subfinder -o output.txt)
|
||
|
||
Args:
|
||
tool_name: 工具名称
|
||
command: 扫描命令(包含输出文件参数)
|
||
timeout: 超时时间(秒)
|
||
log_file: 日志文件路径(可选)
|
||
|
||
Returns:
|
||
执行结果字典(包含 duration 字段)
|
||
|
||
Raises:
|
||
RuntimeError: 执行失败或超时
|
||
"""
|
||
return _executor.execute_and_wait(tool_name, command, timeout, log_file)
|
||
|
||
|
||
def execute_stream(
|
||
cmd: str,
|
||
tool_name: str,
|
||
cwd: Optional[str] = None,
|
||
shell: bool = False,
|
||
encoding: str = 'utf-8',
|
||
suffix_char: Optional[str] = None,
|
||
timeout: Optional[int] = None,
|
||
log_file: Optional[str] = None
|
||
) -> Generator[str, None, None]:
|
||
"""
|
||
流式执行命令(快捷函数)
|
||
|
||
适用场景:工具流式输出 JSON(如 naabu -json)
|
||
|
||
Args:
|
||
cmd: 要执行的命令
|
||
tool_name: 工具名称
|
||
cwd: 工作目录
|
||
shell: 是否使用 shell 执行
|
||
encoding: 编码格式
|
||
suffix_char: 末尾后缀字符
|
||
timeout: 命令执行超时时间(秒)
|
||
log_file: 日志文件路径(可选)
|
||
|
||
Yields:
|
||
str: 每行输出的内容
|
||
|
||
Raises:
|
||
subprocess.TimeoutExpired: 命令执行超时
|
||
"""
|
||
return _executor.execute_stream(cmd, tool_name, cwd, shell, encoding, suffix_char, timeout, log_file)
|