mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-05-18 22:08:13 +02:00
Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 7d7207c12f | |||
| 9eb47d96f5 | |||
| cf1c9c199c | |||
| ce5f20c11e |
+1
-1
@@ -10,7 +10,7 @@
|
|||||||
# ============================================
|
# ============================================
|
||||||
|
|
||||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||||
version: "v1.4.9"
|
version: "v1.4.10"
|
||||||
# 服务器配置
|
# 服务器配置
|
||||||
server:
|
server:
|
||||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
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 |
@@ -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()
|
defer cancel()
|
||||||
|
|
||||||
s.logger.Info("开始执行工具",
|
s.logger.Info("开始执行工具",
|
||||||
|
|||||||
@@ -6,7 +6,9 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"os"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
|
"runtime"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -16,6 +18,7 @@ import (
|
|||||||
"cyberstrike-ai/internal/mcp"
|
"cyberstrike-ai/internal/mcp"
|
||||||
"cyberstrike-ai/internal/storage"
|
"cyberstrike-ai/internal/storage"
|
||||||
|
|
||||||
|
"github.com/creack/pty"
|
||||||
"go.uber.org/zap"
|
"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...)
|
cmd := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
|
||||||
|
applyDefaultTerminalEnv(cmd)
|
||||||
|
|
||||||
e.logger.Info("执行安全工具",
|
e.logger.Info("执行安全工具",
|
||||||
zap.String("tool", toolName),
|
zap.String("tool", toolName),
|
||||||
@@ -160,10 +164,26 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
|
|||||||
// 如果上层提供了 stdout/stderr 增量回调,则边执行边读取并回调。
|
// 如果上层提供了 stdout/stderr 增量回调,则边执行边读取并回调。
|
||||||
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
|
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
|
||||||
output, err = streamCommandOutput(cmd, cb)
|
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 {
|
} else {
|
||||||
outputBytes, err2 := cmd.CombinedOutput()
|
outputBytes, err2 := cmd.CombinedOutput()
|
||||||
output = string(outputBytes)
|
output = string(outputBytes)
|
||||||
err = err2
|
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 {
|
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 {
|
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
|
||||||
output, err = streamCommandOutput(cmd, cb)
|
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 {
|
} else {
|
||||||
outputBytes, err2 := cmd.CombinedOutput()
|
outputBytes, err2 := cmd.CombinedOutput()
|
||||||
output = string(outputBytes)
|
output = string(outputBytes)
|
||||||
err = err2
|
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 {
|
if err != nil {
|
||||||
e.logger.Error("系统命令执行失败",
|
e.logger.Error("系统命令执行失败",
|
||||||
@@ -1066,6 +1104,123 @@ func streamCommandOutput(cmd *exec.Cmd, cb ToolOutputCallback) (string, error) {
|
|||||||
return outBuilder.String(), waitErr
|
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 执行内部工具(不执行外部命令)
|
// executeInternalTool 执行内部工具(不执行外部命令)
|
||||||
func (e *Executor) executeInternalTool(ctx context.Context, toolName string, command string, args map[string]interface{}) (*mcp.ToolResult, error) {
|
func (e *Executor) executeInternalTool(ctx context.Context, toolName string, command string, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||||
// 提取内部工具类型(去掉 "internal:" 前缀)
|
// 提取内部工具类型(去掉 "internal:" 前缀)
|
||||||
|
|||||||
Reference in New Issue
Block a user