From 9eb47d96f5b073993d059b8efbef15848ac8c0a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Wed, 8 Apr 2026 00:18:07 +0800 Subject: [PATCH] Add files via upload --- internal/security/executor.go | 155 ++++++++++++++++++++++++++++++++++ 1 file changed, 155 insertions(+) diff --git a/internal/security/executor.go b/internal/security/executor.go index d56ac46d..70e0dd52 100644 --- a/internal/security/executor.go +++ b/internal/security/executor.go @@ -6,7 +6,9 @@ import ( "encoding/json" "fmt" "io" + "os" "os/exec" + "runtime" "strconv" "strings" "sync" @@ -16,6 +18,7 @@ import ( "cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/storage" + "github.com/creack/pty" "go.uber.org/zap" ) @@ -149,6 +152,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st // 执行命令 cmd := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) + applyDefaultTerminalEnv(cmd) e.logger.Info("执行安全工具", zap.String("tool", toolName), @@ -160,10 +164,26 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st // 如果上层提供了 stdout/stderr 增量回调，则边执行边读取并回调。 if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil { output, err = streamCommandOutput(cmd, cb) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到工具需要 TTY，使用 PTY 重试", + zap.String("tool", toolName), + ) + cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) + applyDefaultTerminalEnv(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, cb) + } } else { outputBytes, err2 := cmd.CombinedOutput() output = string(outputBytes) err = err2 + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到工具需要 TTY，使用 PTY 重试", + zap.String("tool", toolName), + ) + cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) + applyDefaultTerminalEnv(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, nil) + } } if err != nil { // 检查退出码是否在允许列表中 @@ -956,10 +976,28 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int // 若上层提供工具输出增量回调，则边执行边流式读取。 if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil { output, err = streamCommandOutput(cmd, cb) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到系统命令需要 TTY，使用 PTY 重试") + cmd2 := exec.CommandContext(ctx, shell, "-c", command) + if workDir != "" { + cmd2.Dir = workDir + } + applyDefaultTerminalEnv(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, cb) + } } else { outputBytes, err2 := cmd.CombinedOutput() output = string(outputBytes) err = err2 + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到系统命令需要 TTY，使用 PTY 重试") + cmd2 := exec.CommandContext(ctx, shell, "-c", command) + if workDir != "" { + cmd2.Dir = workDir + } + applyDefaultTerminalEnv(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, nil) + } } if err != nil { e.logger.Error("系统命令执行失败", @@ -1066,6 +1104,123 @@ func streamCommandOutput(cmd *exec.Cmd, cb ToolOutputCallback) (string, error) { return outBuilder.String(), waitErr } +// applyDefaultTerminalEnv 为外部工具补齐常见的终端环境变量。 +// 注意：这不会创建 TTY，只是减少某些工具在非交互环境下的“奇怪排版/检测失败”。 +func applyDefaultTerminalEnv(cmd *exec.Cmd) { + if cmd == nil { + return + } + // 仅在未显式设置 Env 时，继承当前进程环境 + if cmd.Env == nil { + cmd.Env = os.Environ() + } + // 如果用户已设置 TERM/COLUMNS/LINES，则不覆盖 + has := func(k string) bool { + prefix := k + "=" + for _, e := range cmd.Env { + if strings.HasPrefix(e, prefix) { + return true + } + } + return false + } + if !has("TERM") { + cmd.Env = append(cmd.Env, "TERM=xterm-256color") + } + if !has("COLUMNS") { + cmd.Env = append(cmd.Env, "COLUMNS=256") + } + if !has("LINES") { + cmd.Env = append(cmd.Env, "LINES=40") + } +} + +func shouldRetryWithPTY(output string) bool { + o := strings.ToLower(output) + // autorecon / python termios 常见报错 + if strings.Contains(o, "inappropriate ioctl for device") { + return true + } + if strings.Contains(o, "termios.error") { + return true + } + // 兜底：stdin 不是 tty + if strings.Contains(o, "not a tty") { + return true + } + return false +} + +// runCommandWithPTY 为子进程分配 PTY，适配需要交互式终端的工具（如 autorecon）。 +// 若 cb != nil，将持续回调增量输出（用于 SSE）。 +func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback) (string, error) { + if runtime.GOOS == "windows" { + // PTY 方案为类 Unix；Windows 走原逻辑 + if cb != nil { + return streamCommandOutput(cmd, cb) + } + out, err := cmd.CombinedOutput() + return string(out), err + } + + ptmx, err := pty.Start(cmd) + if err != nil { + return "", err + } + defer func() { _ = ptmx.Close() }() + + // ctx 取消时尽快终止子进程 + done := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + _ = ptmx.Close() // 触发读退出 + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + case <-done: + } + }() + defer close(done) + + var outBuilder strings.Builder + var deltaBuilder strings.Builder + lastFlush := time.Now() + flush := func() { + if cb == nil || deltaBuilder.Len() == 0 { + deltaBuilder.Reset() + lastFlush = time.Now() + return + } + cb(deltaBuilder.String()) + deltaBuilder.Reset() + lastFlush = time.Now() + } + + buf := make([]byte, 4096) + for { + n, readErr := ptmx.Read(buf) + if n > 0 { + chunk := string(buf[:n]) + // 统一换行为 \n，避免前端错位 + chunk = strings.ReplaceAll(chunk, "\r\n", "\n") + chunk = strings.ReplaceAll(chunk, "\r", "\n") + outBuilder.WriteString(chunk) + deltaBuilder.WriteString(chunk) + if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond { + flush() + } + } + if readErr != nil { + break + } + } + flush() + + waitErr := cmd.Wait() + return outBuilder.String(), waitErr +} + // executeInternalTool 执行内部工具（不执行外部命令） func (e *Executor) executeInternalTool(ctx context.Context, toolName string, command string, args map[string]interface{}) (*mcp.ToolResult, error) { // 提取内部工具类型（去掉 "internal:" 前缀）