CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-06-04 21:38:08 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2026-05-12 12:56:11 +08:00
committed by GitHub
parent 2e1e57ce27
commit cae3a77331
4 changed files with 120 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/security/executor.go
+34 -13
View File
@@ -153,6 +153,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
// 执行命令
cmd := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
applyDefaultTerminalEnv(cmd)
_ = prepareShellCmdSession(cmd)
e.logger.Info("执行安全工具",
zap.String("tool", toolName),
@@ -163,13 +164,14 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
var err error
// 如果上层提供了 stdout/stderr 增量回调，则边执行边读取并回调。
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
output, err = streamCommandOutput(cmd, cb)
output, err = streamCommandOutput(ctx, 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)
_ = prepareShellCmdSession(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, cb)
}
} else {
@@ -182,6 +184,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
)
cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
applyDefaultTerminalEnv(cmd2)
_ = prepareShellCmdSession(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, nil)
}
}
@@ -837,6 +840,8 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
} else {
cmd = exec.CommandContext(ctx, shell, "-c", command)
}
applyDefaultTerminalEnv(cmd)
_ = prepareShellCmdSession(cmd)
// 执行命令
e.logger.Info("执行系统命令",
@@ -865,6 +870,8 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
} else {
pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand)
}
applyDefaultTerminalEnv(pidCmd)
_ = prepareShellCmdSession(pidCmd)
// 获取stdout管道
stdout, err := pidCmd.StdoutPipe()
@@ -976,7 +983,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
var err error
// 若上层提供工具输出增量回调，则边执行边流式读取。
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
output, err = streamCommandOutput(cmd, cb)
output, err = streamCommandOutput(ctx, cmd, cb)
if err != nil && shouldRetryWithPTY(output) {
e.logger.Info("检测到系统命令需要 TTY，使用 PTY 重试")
cmd2 := exec.CommandContext(ctx, shell, "-c", command)
@@ -984,6 +991,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
cmd2.Dir = workDir
}
applyDefaultTerminalEnv(cmd2)
_ = prepareShellCmdSession(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, cb)
}
} else {
@@ -997,6 +1005,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
cmd2.Dir = workDir
}
applyDefaultTerminalEnv(cmd2)
_ = prepareShellCmdSession(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, nil)
}
}
@@ -1034,8 +1043,11 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
}
// streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。
// 保持输出内容完整拼接返回，并用 cb(chunk) 向上层持续推送。
func streamCommandOutput(cmd *exec.Cmd, cb ToolOutputCallback) (string, error) {
// 使用定长块读取，避免按行读取在无换行输出时永久阻塞；ctx 取消时终止进程树。
func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback) (string, error) {
if err := prepareShellCmdSession(cmd); err != nil {
return "", err
}
stdoutPipe, err := cmd.StdoutPipe()
if err != nil {
return "", err
@@ -1051,18 +1063,27 @@ func streamCommandOutput(cmd *exec.Cmd, cb ToolOutputCallback) (string, error) {
return "", err
}
stopWatch := make(chan struct{})
go func() {
select {
case <-ctx.Done():
terminateCmdTree(cmd)
case <-stopWatch:
}
}()
defer close(stopWatch)
chunks := make(chan string, 64)
var wg sync.WaitGroup
readFn := func(r io.Reader) {
defer wg.Done()
br := bufio.NewReader(r)
buf := make([]byte, 8192)
for {
s, readErr := br.ReadString('\n')
if s != "" {
chunks <- s
n, readErr := r.Read(buf)
if n > 0 {
chunks <- string(buf[:n])
}
if readErr != nil {
// EOF 正常结束
return
}
}
@@ -1158,12 +1179,14 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
if runtime.GOOS == "windows" {
// PTY 方案为类 Unix；Windows 走原逻辑
if cb != nil {
return streamCommandOutput(cmd, cb)
return streamCommandOutput(ctx, cmd, cb)
}
_ = prepareShellCmdSession(cmd)
out, err := cmd.CombinedOutput()
return string(out), err
}
_ = prepareShellCmdSession(cmd)
ptmx, err := pty.Start(cmd)
if err != nil {
return "", err
@@ -1176,9 +1199,7 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
select {
case <-ctx.Done():
_ = ptmx.Close() // 触发读退出
if cmd.Process != nil {
_ = cmd.Process.Kill()
}
terminateCmdTree(cmd)
case <-done:
}
}()
internal/security/procattr_unix.go
+31
View File
@@ -0,0 +1,31 @@
//go:build !windows
package security
import (
"os/exec"
"syscall"
)
// prepareShellCmdSession 让 shell 子进程在独立会话中运行，便于超时/取消时整组 SIGKILL（含子进程）。
func prepareShellCmdSession(cmd *exec.Cmd) error {
if cmd == nil {
return nil
}
if cmd.SysProcAttr == nil {
cmd.SysProcAttr = &syscall.SysProcAttr{}
}
cmd.SysProcAttr.Setsid = true
return nil
}
// terminateCmdTree 尽力终止 cmd 及其进程组（Unix 下 Setsid 后 PGID == 首进程 PID）。
func terminateCmdTree(cmd *exec.Cmd) {
if cmd == nil || cmd.Process == nil {
return
}
pid := cmd.Process.Pid
if err := syscall.Kill(-pid, syscall.SIGKILL); err != nil {
_ = cmd.Process.Kill()
}
}
internal/security/procattr_windows.go
+17
View File
@@ -0,0 +1,17 @@
//go:build windows
package security
import "os/exec"
func prepareShellCmdSession(cmd *exec.Cmd) error {
_ = cmd
return nil
}
func terminateCmdTree(cmd *exec.Cmd) {
if cmd == nil || cmd.Process == nil {
return
}
_ = cmd.Process.Kill()
}