Logan.Li
e0d652d211
- 支持 Windows Agent 构建、安装脚本和统一上报协议 - 增强 RDP 剪贴板、虚拟盘和文件传输能力 - 新增 AIOps 上下文工具、告警 RCA trace 落库和回放接口 - 重构 docs 模块文档并补充终端、堡垒机、监控、AI Agent 设计
104 lines
2.6 KiB
Go
104 lines
2.6 KiB
Go
//go:build windows
|
||
|
||
// Package task 实现 Windows Agent 的 PowerShell 任务执行能力。 // Logan.Li
|
||
package task
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"fmt"
|
||
"log"
|
||
"os"
|
||
"os/exec"
|
||
"path/filepath"
|
||
"time"
|
||
)
|
||
|
||
type Executor struct {
|
||
logger *log.Logger
|
||
}
|
||
|
||
func New(logger *log.Logger) *Executor {
|
||
return &Executor{logger: logger}
|
||
}
|
||
|
||
func (e *Executor) ExecuteTask(ctx context.Context, t Task) Result {
|
||
res := Result{TaskID: t.TaskID}
|
||
execCtx, cancel := e.withTimeout(ctx, t.Timeout)
|
||
defer cancel()
|
||
|
||
var cmd *exec.Cmd
|
||
var cleanup func()
|
||
switch t.Type {
|
||
case TypeCommand:
|
||
cmd = exec.CommandContext(execCtx, "powershell.exe", "-NoProfile", "-ExecutionPolicy", "Bypass", "-Command", t.Command)
|
||
case TypeScript:
|
||
path, err := writeTempPowerShell(t.Command)
|
||
if err != nil {
|
||
res.ExitCode = -1
|
||
res.Status = "failed"
|
||
res.Result = "写入临时 PowerShell 脚本失败: " + err.Error()
|
||
return res
|
||
}
|
||
cleanup = func() { _ = os.Remove(path) }
|
||
cmd = exec.CommandContext(execCtx, "powershell.exe", "-NoProfile", "-ExecutionPolicy", "Bypass", "-File", path)
|
||
default:
|
||
res.ExitCode = -1
|
||
res.Status = "failed"
|
||
res.Result = fmt.Sprintf("未知的任务类型: %q", t.Type)
|
||
return res
|
||
}
|
||
if cleanup != nil {
|
||
defer cleanup()
|
||
}
|
||
|
||
var buf bytes.Buffer
|
||
cmd.Stdout = &buf
|
||
cmd.Stderr = &buf
|
||
start := time.Now()
|
||
err := cmd.Run()
|
||
elapsed := time.Since(start)
|
||
if execCtx.Err() == context.DeadlineExceeded {
|
||
res.ExitCode = -1
|
||
res.Status = "timeout"
|
||
res.Result = fmt.Sprintf("任务执行超时（%ds），已输出:\n%s", t.Timeout, buf.String())
|
||
e.logger.Printf("任务 %s 执行超时（%ds）", t.TaskID, t.Timeout)
|
||
return res
|
||
}
|
||
res.ExitCode = exitCode(err)
|
||
res.Result = buf.String()
|
||
if res.ExitCode == 0 {
|
||
res.Status = "success"
|
||
e.logger.Printf("任务 %s 执行成功 exit=0 耗时=%s 输出=%dB", t.TaskID, elapsed, buf.Len())
|
||
} else {
|
||
res.Status = "failed"
|
||
e.logger.Printf("任务 %s 执行失败 exit=%d 耗时=%s 输出=%dB", t.TaskID, res.ExitCode, elapsed, buf.Len())
|
||
}
|
||
return res
|
||
}
|
||
|
||
func (e *Executor) withTimeout(parent context.Context, timeout int) (context.Context, context.CancelFunc) {
|
||
if timeout <= 0 {
|
||
return parent, func() {}
|
||
}
|
||
return context.WithTimeout(parent, time.Duration(timeout)*time.Second)
|
||
}
|
||
|
||
func exitCode(err error) int {
|
||
if err == nil {
|
||
return 0
|
||
}
|
||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||
return exitErr.ExitCode()
|
||
}
|
||
return -1
|
||
}
|
||
|
||
func writeTempPowerShell(script string) (string, error) {
|
||
path := filepath.Join(os.TempDir(), fmt.Sprintf("aiops-agent-task-%d.ps1", time.Now().UnixNano()))
|
||
if err := os.WriteFile(path, []byte(script), 0600); err != nil {
|
||
return "", err
|
||
}
|
||
return path, nil
|
||
}
|