Compare commits

..

4 Commits

Author SHA1 Message Date
公明 7d7207c12f Update config.yaml 2026-04-08 16:58:20 +08:00
公明 9eb47d96f5 Add files via upload 2026-04-08 00:18:07 +08:00
公明 cf1c9c199c Update server.go 2026-04-07 11:51:35 +08:00
公明 ce5f20c11e Add files via upload 2026-04-04 15:05:38 +08:00
4 changed files with 157 additions and 2 deletions
+1 -1
View File
@@ -10,7 +10,7 @@
# ============================================
# 前端显示的版本号(可选,不填则显示默认版本)
version: "v1.4.9"
version: "v1.4.10"
# 服务器配置
server:
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
Binary file not shown.

Before

Width:  |  Height:  |  Size: 178 KiB

After

Width:  |  Height:  |  Size: 182 KiB

+1 -1
View File
@@ -444,7 +444,7 @@ func (s *Server) handleCallTool(msg *Message) *Message {
}
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
defer cancel()
s.logger.Info("开始执行工具",
+155
View File
@@ -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 方案为类 UnixWindows 走原逻辑
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:" 前缀)