359 lines
12 KiB
Go
359 lines
12 KiB
Go
package script
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"os"
|
||
"os/exec"
|
||
"path/filepath"
|
||
"strings"
|
||
"syscall"
|
||
"time"
|
||
|
||
"github.com/infraboard/mcube/v2/ioc"
|
||
"github.com/infraboard/mcube/v2/ioc/config/log"
|
||
"github.com/rs/zerolog"
|
||
)
|
||
|
||
// 把ScriptExcutor 托管给IOC管理
|
||
func init() {
|
||
ioc.Controller().Registry(&ScriptExcutor{
|
||
WorkDirPrefix: "workspace",
|
||
ScriptDirPrefix: "shells",
|
||
IsVerifyScriptIntegrity: false,
|
||
runningCommands: make(map[string]*ExecuteScriptRequest),
|
||
})
|
||
}
|
||
|
||
// ScriptExcutor 脚本执行器，负责制定脚本的执行逻辑，比如准备环境变量、执行脚本、处理输出等
|
||
type ScriptExcutor struct {
|
||
ioc.ObjectImpl
|
||
|
||
// 工作目录前缀，用于确保脚本执行在指定的目录下
|
||
WorkDirPrefix string `toml:"work_dir_prefix" json:"work_dir_prefix" yaml:"work_dir_prefix"`
|
||
// 脚本目录前缀，用于确保脚本路径在指定的目录下, 避免执行不安全的脚本
|
||
ScriptDirPrefix string `toml:"script_dir_prefix" json:"script_dir_prefix" yaml:"script_dir_prefix"`
|
||
// 是否启用脚本完整性校验, 通过编译时嵌入的脚本 hash 来校验脚本的完整性, 防止脚本被篡改
|
||
IsVerifyScriptIntegrity bool `toml:"is_verify_script_integrity" json:"is_verify_script_integrity" yaml:"is_verify_script_integrity"`
|
||
|
||
// 日志记录器，用于记录脚本执行的日志信息
|
||
log *zerolog.Logger
|
||
// 脚本完整性管理器，用于校验脚本的完整性
|
||
integrityManager *ScriptIntegrityManager
|
||
|
||
// 当前正在执行中的命令列表，key为工作目录，value为执行请求
|
||
runningCommands map[string]*ExecuteScriptRequest
|
||
}
|
||
|
||
func (e *ScriptExcutor) Init() error {
|
||
// 初始化日志记录器
|
||
e.log = log.Sub(e.Name())
|
||
// 初始化脚本完整性管理器
|
||
e.integrityManager = NewScriptIntegrityManager(e.ScriptDirPrefix, e.IsVerifyScriptIntegrity)
|
||
return nil
|
||
}
|
||
|
||
func (e *ScriptExcutor) Name() string {
|
||
return APP_NAME
|
||
}
|
||
|
||
// ExecuteScript 执行脚本的核心方法，接受一个 ExecuteScriptRequest 请求，返回一个 ExecutionResult 结果
|
||
func (e *ScriptExcutor) ExecuteScript(ctx context.Context, in *ExecuteScriptRequest) (*ExecutionResult, error) {
|
||
// 确保工作目录前缀, 避免脚本执行在不安全的目录下, 同时也方便清理执行结果
|
||
in = in.WithWorkspacePrefix(e.WorkDirPrefix)
|
||
|
||
// 确保工作目录存在
|
||
if err := e.ensureWorkDir(in.workDir); err != nil {
|
||
e.log.Error().Str("task_id", in.metadata.ID).Str("work_dir", in.workDir).Err(err).Msg("创建工作目录失败")
|
||
return nil, err
|
||
}
|
||
|
||
// 确保元数据存在
|
||
if in.metadata == nil {
|
||
in.metadata = &CommandMetadata{
|
||
ID: generateID(),
|
||
CreatedBy: getCurrentUser(),
|
||
CreatedAt: time.Now(),
|
||
WorkDir: in.workDir,
|
||
EnvVars: in.envVars,
|
||
Timeout: in.timeout,
|
||
}
|
||
}
|
||
|
||
// zerolog 的用法
|
||
e.log.Info().Str("task_id", in.metadata.ID).Str("script_path", in.ScriptPath).Msg("准备执行脚本")
|
||
|
||
// 脚本路径消毒
|
||
scriptPath := e.WithPrefixScripPath(in.ScriptPath)
|
||
fullScriptPath, err := e.sanitizeScriptPath(scriptPath)
|
||
if err != nil {
|
||
e.log.Error().Str("task_id", in.metadata.ID).Str("script_path", scriptPath).Err(err).Msg("脚本路径验证失败")
|
||
return nil, err
|
||
}
|
||
e.log.Info().Str("task_id", in.metadata.ID).Str("full_script_path", fullScriptPath).Msg("脚本路径验证成功")
|
||
|
||
// 校验脚本完整性
|
||
if err := e.integrityManager.VerifyScript(fullScriptPath); err != nil {
|
||
e.log.Error().Str("task_id", in.metadata.ID).Str("script_path", fullScriptPath).Err(err).Msg("脚本完整性校验失败")
|
||
return nil, fmt.Errorf("脚本完整性校验失败: %v", err)
|
||
}
|
||
|
||
// 构建完整的shell命令
|
||
shellArgs := []string{fullScriptPath}
|
||
shellArgs = append(shellArgs, in.Args...)
|
||
shellCommand := strings.Join(shellArgs, " ")
|
||
|
||
e.log.Info().Str("task_id", in.metadata.ID).Str("command", shellCommand).Msg("脚本执行命令")
|
||
|
||
// 注入标准变量
|
||
e.InjectEnv(in)
|
||
|
||
// 根据超时设置创建 context
|
||
execCtx := ctx
|
||
var cancel context.CancelFunc
|
||
if in.timeout > 0 {
|
||
execCtx, cancel = context.WithTimeout(ctx, in.timeout)
|
||
defer cancel()
|
||
e.log.Info().Str("task_id", in.metadata.ID).Dur("timeout", in.timeout).Msg("设置脚本执行超时时间")
|
||
}
|
||
|
||
// 创建命令
|
||
cmd := exec.CommandContext(execCtx, "/bin/sh", "-c", shellCommand)
|
||
cmd.Dir = in.workDir
|
||
cmd.Env = in.buildEnv()
|
||
// 根据参数设置进程组属性
|
||
if in.createProcessGroup {
|
||
cmd.SysProcAttr = &syscall.SysProcAttr{
|
||
Setpgid: true, // 创建新的进程组
|
||
}
|
||
}
|
||
|
||
// 根据参数设置自定义 Cancel 函数
|
||
if in.useProcessGroupKill {
|
||
// 自定义 Cancel 函数，确保杀死进程组
|
||
cmd.Cancel = func() error {
|
||
if cmd.Process == nil {
|
||
return nil
|
||
}
|
||
// 使用负 PID 杀死整个进程组
|
||
return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
|
||
}
|
||
}
|
||
|
||
in.cmd = cmd
|
||
|
||
// 脚本执行生成调试脚本（避免被脚本内部的 git clone 等操作清空）
|
||
if in.isDebugScriptEnabled() {
|
||
if debugScriptPath, writeErr := in.WriteDebugScript(fullScriptPath, in.Args); writeErr != nil {
|
||
e.log.Warn().Str("task_id", in.metadata.ID).Err(writeErr).Msg("生成调试脚本失败")
|
||
} else if debugScriptPath != "" {
|
||
e.log.Info().Str("task_id", in.metadata.ID).Str("debug_script", debugScriptPath).Msg("已生成调试脚本，可进入工作目录执行 ./debug.sh 进行调试")
|
||
}
|
||
}
|
||
|
||
// 执行脚本，并且把结果写到输出到文件里面
|
||
// 开始执行命令
|
||
|
||
// 添加到运行中的命令列表中
|
||
e.runningCommands[in.workDir] = in
|
||
defer delete(e.runningCommands, in.workDir)
|
||
|
||
taskId := "unknown"
|
||
if in.metadata != nil {
|
||
taskId = in.metadata.ID
|
||
}
|
||
|
||
// 获取日志写入器
|
||
logWriter, err := in.getLogWriter()
|
||
if err != nil {
|
||
e.log.Error().Str("task_id", taskId).Err(err).Msg("创建日志写入器失败")
|
||
return nil, err
|
||
}
|
||
defer func() {
|
||
if closer, ok := logWriter.(io.Closer); ok {
|
||
closer.Close()
|
||
}
|
||
}()
|
||
|
||
// 命令的日志输出到制定的地方
|
||
in.cmd.Stdout = logWriter
|
||
in.cmd.Stderr = logWriter
|
||
|
||
// 记录执行开始, 包括时间、工作目录、脚本路径和元数据等信息
|
||
startTime := time.Now()
|
||
logHeader := fmt.Sprintf("\n=== Execution Started ===\nTime: %s\nWorkDir: %s\nScript: %s\nMetadata: %s\n",
|
||
startTime.Format("2006-01-02 15:04:05"),
|
||
in.workDir,
|
||
fullScriptPath,
|
||
in.metadata.String(),
|
||
)
|
||
fmt.Fprint(logWriter, logHeader)
|
||
fmt.Fprint(logWriter, "\n")
|
||
|
||
// 执行命令
|
||
e.log.Info().Str("task_id", taskId).Msg("开始执行命令")
|
||
err = in.cmd.Run()
|
||
endTime := time.Now()
|
||
duration := endTime.Sub(startTime)
|
||
|
||
// 构建执行结果
|
||
result := &ExecutionResult{
|
||
Command: fullScriptPath,
|
||
StartTime: startTime,
|
||
EndTime: &endTime,
|
||
Duration: duration,
|
||
Success: err == nil,
|
||
Metadata: in.metadata,
|
||
}
|
||
|
||
if err != nil {
|
||
// 判断是否为超时错误
|
||
if in.cmd.ProcessState != nil && in.cmd.ProcessState.ExitCode() == -1 {
|
||
e.log.Error().Str("task_id", taskId).Dur("duration", duration).Dur("timeout", in.timeout).Err(err).Msg("命令执行超时")
|
||
} else {
|
||
e.log.Error().Str("task_id", taskId).Dur("duration", duration).Err(err).Msg("命令执行失败")
|
||
}
|
||
|
||
// 错误信息
|
||
result.Error = err.Error()
|
||
if in.cmd.ProcessState != nil {
|
||
result.ExitCode = in.cmd.ProcessState.ExitCode()
|
||
} else {
|
||
result.ExitCode = -1
|
||
}
|
||
} else {
|
||
e.log.Info().Str("task_id", taskId).Dur("duration", duration).Msg("命令执行成功")
|
||
}
|
||
|
||
// 解析脚本输出参数（无论成功或失败都收集）
|
||
outputFile := filepath.Join(in.workDir, "output.env")
|
||
if outputParams, parseErr := ParseOutputParams(outputFile); parseErr == nil && len(outputParams) > 0 {
|
||
result.OutputParams = outputParams
|
||
e.log.Info().Str("task_id", taskId).Bool("success", result.Success).Int("param_count", len(outputParams)).Msg("成功解析脚本输出参数")
|
||
} else if parseErr != nil && !os.IsNotExist(parseErr) {
|
||
e.log.Warn().Str("task_id", taskId).Str("output_file", outputFile).Err(parseErr).Msg("解析输出参数文件失败")
|
||
}
|
||
|
||
// 保存执行结果(无论成功或失败都保存结果，方便后续查询和调试)
|
||
e.log.Debug().Str("task_id", taskId).Str("result_file", in.GetResultFilePath()).Msg("准备保存执行结果")
|
||
if saveErr := in.saveResult(result); saveErr != nil {
|
||
e.log.Error().Str("task_id", taskId).Err(saveErr).Msg("保存执行结果失败")
|
||
}
|
||
|
||
// 记录执行结束
|
||
logFooter := fmt.Sprintf("\n=== Execution Finished ===\nTime: %s\nDuration: %v\nSuccess: %t\nExitCode: %d\nError: %v\n",
|
||
endTime.Format("2006-01-02 15:04:05"),
|
||
duration,
|
||
result.Success,
|
||
result.ExitCode,
|
||
err,
|
||
)
|
||
fmt.Fprint(logWriter, logFooter)
|
||
return result, err
|
||
}
|
||
|
||
func (s *ScriptExcutor) GetWorkDirAbsPath(workDir string) (string, error) {
|
||
// 如果是绝对路径，直接返回；如果是相对路径，基于工作目录前缀构建绝对路径
|
||
if filepath.IsAbs(workDir) {
|
||
return workDir, nil
|
||
}
|
||
return filepath.Abs(filepath.Join(s.WorkDirPrefix, workDir))
|
||
}
|
||
|
||
// saveResult 保存执行结果到文件
|
||
func (s *ExecuteScriptRequest) saveResult(result *ExecutionResult) error {
|
||
if result == nil {
|
||
return nil
|
||
}
|
||
|
||
resultPath := s.resultFile
|
||
if !filepath.IsAbs(resultPath) {
|
||
resultPath = filepath.Join(s.workDir, resultPath)
|
||
}
|
||
|
||
// 美化输出JSON
|
||
data, err := json.MarshalIndent(result, "", " ")
|
||
if err != nil {
|
||
return fmt.Errorf("failed to marshal result: %v", err)
|
||
}
|
||
|
||
if err := os.WriteFile(resultPath, data, 0644); err != nil {
|
||
return fmt.Errorf("failed to save result file: %v", err)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (e *ScriptExcutor) InjectEnv(in *ExecuteScriptRequest) {
|
||
// 注入标准变量
|
||
in.SetEnv("SCRIPT_DIR", e.ScriptDirPrefix)
|
||
in.SetEnv("WORKSPACE", in.workDir)
|
||
|
||
// 创建 output.env 文件（用于脚本输出参数）
|
||
outputEnvFile := filepath.Join(in.workDir, "output.env")
|
||
// 创建空文件（如果已存在则清空）
|
||
if f, err := os.Create(outputEnvFile); err == nil {
|
||
f.Close()
|
||
// 注入环境变量，让脚本可以直接使用
|
||
in.SetEnv("OUTPUT_ENV_FILE", outputEnvFile)
|
||
}
|
||
}
|
||
|
||
// ensureWorkDir 确保工作目录存在
|
||
func (e *ScriptExcutor) ensureWorkDir(workspace string) error {
|
||
if workspace == "" {
|
||
return fmt.Errorf("work directory cannot be empty")
|
||
}
|
||
|
||
if _, err := os.Stat(workspace); os.IsNotExist(err) {
|
||
if err := os.MkdirAll(workspace, 0755); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (r *ScriptExcutor) WithPrefixScripPath(scriptName string) string {
|
||
return r.ScriptDirPrefix + "/" + scriptName
|
||
}
|
||
|
||
// sanitizeScriptPath 脚本路径消毒, 确保脚本路径在指定的目录下，并且文件存在且可执行
|
||
func (s *ScriptExcutor) sanitizeScriptPath(scriptPath string) (string, error) {
|
||
if strings.TrimSpace(scriptPath) == "" {
|
||
return "", fmt.Errorf("script path cannot be empty")
|
||
}
|
||
|
||
// 确保脚本路径在工作目录内
|
||
scriptFullPath := scriptPath
|
||
if !filepath.IsAbs(scriptPath) {
|
||
// 直接获取绝对路径
|
||
fullPath, err := filepath.Abs(scriptPath)
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to get absolute path for script path %s: %v", scriptFullPath, err)
|
||
}
|
||
scriptFullPath = fullPath
|
||
}
|
||
|
||
// 检查文件是否存在
|
||
if _, err := os.Stat(scriptFullPath); os.IsNotExist(err) {
|
||
return "", fmt.Errorf("script file does not exist: %s", scriptFullPath)
|
||
}
|
||
|
||
// 检查文件是否可执行
|
||
info, err := os.Stat(scriptFullPath)
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to get script file info: %v", err)
|
||
}
|
||
|
||
// 如果是脚本文件，确保它有执行权限
|
||
if info.Mode().Perm()&0111 == 0 {
|
||
// 尝试添加执行权限
|
||
if err := os.Chmod(scriptFullPath, info.Mode()|0111); err != nil {
|
||
return "", fmt.Errorf("script file is not executable and cannot be made executable: %v", err)
|
||
}
|
||
}
|
||
|
||
return scriptFullPath, nil
|
||
}
|