ops-assistant/internal/core/runbook/executor.go

package runbook

import (
	"archive/tar"
	"archive/zip"
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"ops-assistant/internal/core/ecode"
	"ops-assistant/models"

	"gorm.io/gorm"
)

type Executor struct {
	db         *gorm.DB
	runbookDir string
}

func NewExecutor(db *gorm.DB, runbookDir string) *Executor {
	return &Executor{db: db, runbookDir: runbookDir}
}

func (e *Executor) Run(commandText, runbookName string, operator int64) (uint, string, error) {
	return e.RunWithInputsAndMeta(commandText, runbookName, operator, map[string]string{}, NewMeta())
}

func (e *Executor) RunWithInputs(commandText, runbookName string, operator int64, inputs map[string]string) (uint, string, error) {
	return e.RunWithInputsAndMeta(commandText, runbookName, operator, inputs, NewMeta())
}

func (e *Executor) RunWithInputsAndMeta(commandText, runbookName string, operator int64, inputs map[string]string, meta RunMeta) (uint, string, error) {
	started := time.Now()
	inputJSON := "{}"
	if b, err := json.Marshal(inputs); err == nil {
		inputJSON = string(b)
	}
	job := models.OpsJob{
		Command:     commandText,
		Runbook:     runbookName,
		Operator:    operator,
		Target:      strings.TrimSpace(meta.Target),
		RiskLevel:   strings.TrimSpace(meta.RiskLevel),
		RequestID:   strings.TrimSpace(meta.RequestID),
		ConfirmHash: strings.TrimSpace(meta.ConfirmHash),
		InputJSON:   inputJSON,
		Status:      "pending",
		StartedAt:   started,
	}
	if job.RiskLevel == "" {
		job.RiskLevel = "low"
	}
	if err := e.db.Create(&job).Error; err != nil {
		return 0, "", err
	}

	release := acquireTargetLock(job.Target)
	defer release()

	job.Status = "running"
	_ = e.db.Save(&job).Error

	specPath := filepath.Join(e.runbookDir, runbookName+".yaml")
	data, err := os.ReadFile(specPath)
	if err != nil {
		e.finishJob(&job, "failed", "runbook not found")
		return job.ID, "", err
	}
	spec, err := Parse(data)
	if err != nil {
		e.finishJob(&job, "failed", "runbook parse failed")
		return job.ID, "", err
	}

	outputs := map[string]string{}
	ctx := map[string]string{}
	dataOut := "ok"

	jobCtx, jobCancel := context.WithCancel(context.Background())
	registerJobCancel(job.ID, jobCancel)
	defer clearJobCancel(job.ID)
	for k, v := range inputs {
		ctx["inputs."+k] = v
	}
	if t := strings.TrimSpace(os.Getenv("CPA_MANAGEMENT_BASE")); t != "" {
		ctx["env.cpa_management_base"] = t
	} else {
		var sset models.AppSetting
		if err := e.db.Where("key = ?", "cpa_management_base").First(&sset).Error; err == nil {
			if strings.TrimSpace(sset.Value) != "" {
				ctx["env.cpa_management_base"] = strings.TrimSpace(sset.Value)
			}
		}
	}
	if t := strings.TrimSpace(os.Getenv("CPA_MANAGEMENT_TOKEN")); t != "" {
		ctx["env.cpa_management_token"] = t
	} else {
		var sset models.AppSetting
		if err := e.db.Where("key = ?", "cpa_management_token").First(&sset).Error; err == nil {
			if strings.TrimSpace(sset.Value) != "" {
				ctx["env.cpa_management_token"] = strings.TrimSpace(sset.Value)
			}
		}
	}
	// Cloudflare settings
	if t := strings.TrimSpace(os.Getenv("CF_ACCOUNT_ID")); t != "" {
		ctx["env_cf_account_id"] = t
	} else {
		var sset models.AppSetting
		if err := e.db.Where("key = ?", "cf_account_id").First(&sset).Error; err == nil {
			if strings.TrimSpace(sset.Value) != "" {
				ctx["env_cf_account_id"] = strings.TrimSpace(sset.Value)
			}
		}
	}
	if t := strings.TrimSpace(os.Getenv("CF_API_EMAIL")); t != "" {
		ctx["env_cf_api_email"] = t
	} else {
		var sset models.AppSetting
		if err := e.db.Where("key = ?", "cf_api_email").First(&sset).Error; err == nil {
			if strings.TrimSpace(sset.Value) != "" {
				ctx["env_cf_api_email"] = strings.TrimSpace(sset.Value)
			}
		}
	}
	if t := strings.TrimSpace(os.Getenv("CF_API_TOKEN")); t != "" {
		ctx["env_cf_api_token"] = t
	} else {
		var sset models.AppSetting
		if err := e.db.Where("key = ?", "cf_api_token").First(&sset).Error; err == nil {
			if strings.TrimSpace(sset.Value) != "" {
				ctx["env_cf_api_token"] = strings.TrimSpace(sset.Value)
			}
		}
	}

	// inject input env vars for runbook steps
	for k, v := range inputs {
		if strings.TrimSpace(v) != "" {
			ctx["env.INPUT_"+strings.ToUpper(k)] = v
		}
	}
	for _, st := range spec.Steps {
		if isJobCancelled(e.db, job.ID) {
			e.finishJob(&job, "cancelled", ecode.Tag(ecode.ErrJobCancelled, "cancelled by user"))
			return job.ID, "", fmt.Errorf(ecode.Tag(ecode.ErrJobCancelled, "cancelled by user"))
		}

		rendered := renderStep(st, ctx)
		step := models.OpsJobStep{JobID: job.ID, StepID: rendered.ID, Action: rendered.Action, Status: "running", StartedAt: time.Now()}
		_ = e.db.Create(&step).Error

		timeout := meta.timeoutOrDefault()
		rc, stdout, stderr, runErr := e.execStep(jobCtx, rendered, outputs, timeout)
		step.RC = rc
		step.StdoutTail = tail(stdout, 1200)
		step.StderrTail = tail(stderr, 1200)
		step.EndedAt = time.Now()
		if runErr != nil || rc != 0 {
			step.Status = "failed"
			_ = e.db.Save(&step).Error
			e.finishJob(&job, "failed", fmt.Sprintf("%s: step=%s failed", ecode.ErrStepFailed, rendered.ID))
			if runErr == nil {
				runErr = fmt.Errorf("rc=%d", rc)
			}
			return job.ID, "", fmt.Errorf(ecode.Tag(ecode.ErrStepFailed, fmt.Sprintf("step %s failed: %v", rendered.ID, runErr)))
		}
		step.Status = "success"
		_ = e.db.Save(&step).Error
		outputs[rendered.ID] = stdout
		ctx["steps."+rendered.ID+".output"] = stdout
		if strings.TrimSpace(stdout) != "" && rendered.Action != "assert.json" {
			dataOut = stdout
		}
	}

	e.finishJob(&job, "success", "ok")
	return job.ID, dataOut, nil
}

func (e *Executor) execStep(parent context.Context, st Step, outputs map[string]string, timeout time.Duration) (int, string, string, error) {
	switch st.Action {
	case "ssh.exec":
		target := strings.TrimSpace(fmt.Sprintf("%v", st.With["target"]))
		cmdText := strings.TrimSpace(fmt.Sprintf("%v", st.With["command"]))
		if target == "" || cmdText == "" {
			return 1, "", "missing target/command", fmt.Errorf("missing target/command")
		}
		resolved := resolveTarget(e.db, target)
		if !resolved.Found {
			return 1, "", "invalid target", fmt.Errorf("invalid target: %s", target)
		}
		ctx, cancel := context.WithTimeout(parent, timeout)
		defer cancel()
		args := []string{"-p", strconv.Itoa(resolved.Port), resolved.User + "@" + resolved.Host, cmdText}
		cmd := exec.CommandContext(ctx, "ssh", args...)
		var outb, errb bytes.Buffer
		cmd.Stdout = &outb
		cmd.Stderr = &errb
		err := cmd.Run()
		rc := 0
		if err != nil {
			rc = 1
			if ctx.Err() == context.DeadlineExceeded {
				return rc, strings.TrimSpace(outb.String()), strings.TrimSpace(errb.String()), fmt.Errorf(ecode.Tag(ecode.ErrStepTimeout, "ssh step timeout"))
			}
		}
		return rc, strings.TrimSpace(outb.String()), strings.TrimSpace(errb.String()), err

	case "shell.exec":
		cmdText := strings.TrimSpace(fmt.Sprintf("%v", st.With["command"]))
		if cmdText == "" {
			return 1, "", "missing command", fmt.Errorf("missing command")
		}
		ctx, cancel := context.WithTimeout(parent, timeout)
		defer cancel()
		cmd := exec.CommandContext(ctx, "bash", "-lc", cmdText)
		var outb, errb bytes.Buffer
		cmd.Stdout = &outb
		cmd.Stderr = &errb
		err := cmd.Run()
		rc := 0
		if err != nil {
			rc = 1
			if ctx.Err() == context.DeadlineExceeded {
				return rc, strings.TrimSpace(outb.String()), strings.TrimSpace(errb.String()), fmt.Errorf(ecode.Tag(ecode.ErrStepTimeout, "shell step timeout"))
			}
		}
		return rc, strings.TrimSpace(outb.String()), strings.TrimSpace(errb.String()), err

	case "http.get":
		url := strings.TrimSpace(fmt.Sprintf("%v", st.With["url"]))
		if url == "" {
			return 1, "", "missing url", fmt.Errorf("missing url")
		}
		ctx, cancel := context.WithTimeout(parent, timeout)
		defer cancel()
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
		if err != nil {
			return 1, "", err.Error(), err
		}
		applyHTTPHeaders(req, st.With["headers"])
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			if ctx.Err() == context.DeadlineExceeded {
				return 1, "", "http step timeout", fmt.Errorf(ecode.Tag(ecode.ErrStepTimeout, "http step timeout"))
			}
			return 1, "", err.Error(), err
		}
		defer resp.Body.Close()
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
			return 1, strings.TrimSpace(string(body)), resp.Status, fmt.Errorf("http status: %d", resp.StatusCode)
		}
		return 0, strings.TrimSpace(string(body)), "", nil

	case "http.post":
		url := strings.TrimSpace(fmt.Sprintf("%v", st.With["url"]))
		if url == "" {
			return 1, "", "missing url", fmt.Errorf("missing url")
		}
		bodyStr := ""
		if b, ok := st.With["body"]; ok {
			switch t := b.(type) {
			case string:
				bodyStr = t
			default:
				if jb, err := json.Marshal(t); err == nil {
					bodyStr = string(jb)
				}
			}
		}
		ctx, cancel := context.WithTimeout(parent, timeout)
		defer cancel()
		req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(bodyStr))
		if err != nil {
			return 1, "", err.Error(), err
		}
		hasCT := applyHTTPHeaders(req, st.With["headers"])
		if bodyStr != "" && !hasCT {
			req.Header.Set("Content-Type", "application/json")
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			if ctx.Err() == context.DeadlineExceeded {
				return 1, "", "http step timeout", fmt.Errorf(ecode.Tag(ecode.ErrStepTimeout, "http step timeout"))
			}
			return 1, "", err.Error(), err
		}
		defer resp.Body.Close()
		respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
			return 1, strings.TrimSpace(string(respBody)), resp.Status, fmt.Errorf("http status: %d", resp.StatusCode)
		}
		return 0, strings.TrimSpace(string(respBody)), "", nil

	case "file.upload":
		source := strings.TrimSpace(fmt.Sprintf("%v", st.With["source"]))
		stagingDir := strings.TrimSpace(fmt.Sprintf("%v", st.With["staging_dir"]))
		filename := strings.TrimSpace(fmt.Sprintf("%v", st.With["filename"]))
		if source == "" {
			return 1, "", "missing source", fmt.Errorf("missing source")
		}
		if stagingDir == "" {
			stagingDir = "/tmp/ops-staging"
		}
		if filename == "" {
			filename = filepath.Base(source)
		}
		if err := os.MkdirAll(stagingDir, 0o755); err != nil {
			return 1, "", err.Error(), err
		}
		dest := filepath.Join(stagingDir, filepath.Base(filename))
		if err := copyFile(source, dest); err != nil {
			return 1, "", err.Error(), err
		}
		return 0, dest, "", nil

	case "file.extract":
		archivePath := strings.TrimSpace(fmt.Sprintf("%v", st.With["archive"]))
		toDir := strings.TrimSpace(fmt.Sprintf("%v", st.With["to_dir"]))
		if archivePath == "" {
			return 1, "", "missing archive", fmt.Errorf("missing archive")
		}
		if toDir == "" {
			toDir = "/tmp/ops-staging/extracted"
		}
		if err := os.MkdirAll(toDir, 0o755); err != nil {
			return 1, "", err.Error(), err
		}
		lower := strings.ToLower(archivePath)
		switch {
		case strings.HasSuffix(lower, ".zip"):
			if err := extractZip(archivePath, toDir); err != nil {
				return 1, "", err.Error(), err
			}
		case strings.HasSuffix(lower, ".tar.gz") || strings.HasSuffix(lower, ".tgz"):
			if err := extractTarGz(archivePath, toDir); err != nil {
				return 1, "", err.Error(), err
			}
		case strings.HasSuffix(lower, ".tar"):
			if err := extractTar(archivePath, toDir); err != nil {
				return 1, "", err.Error(), err
			}
		default:
			return 1, "", "unsupported archive format", fmt.Errorf("unsupported archive format: %s", archivePath)
		}
		return 0, toDir, "", nil

	case "assert.json":
		sourceStep := strings.TrimSpace(fmt.Sprintf("%v", st.With["source_step"]))
		if sourceStep == "" {
			return 1, "", "missing source_step", fmt.Errorf("missing source_step")
		}
		raw, ok := outputs[sourceStep]
		if !ok {
			return 1, "", "source step output not found", fmt.Errorf("source step output not found: %s", sourceStep)
		}

		var payload any
		if err := json.Unmarshal([]byte(raw), &payload); err != nil {
			return 1, "", "invalid json", err
		}

		expectSuccess := parseExpectSuccess(st.With["expect_success"])
		if expectSuccess {
			v, ok := lookupPath(payload, "success")
			if !ok {
				return 1, "", "json path not found: success", fmt.Errorf("json path not found: success")
			}
			if !truthyValue(v) {
				errSummary := summarizeJSONErrors(payload)
				return 1, "", "api success=false: " + errSummary, fmt.Errorf("api success=false: %s", errSummary)
			}
		}

		rules := parseRequiredPaths(st.With["required_paths"])
		for _, p := range rules {
			if _, ok := lookupPath(payload, p); !ok {
				return 1, "", "json path not found: " + p, fmt.Errorf("json path not found: %s", p)
			}
		}
		if len(rules) == 0 && !expectSuccess {
			return 1, "", "required_paths empty", fmt.Errorf("required_paths empty")
		}
		return 0, "assert ok", "", nil

	case "sleep":
		ms := 1000
		if v, ok := st.With["ms"]; ok {
			switch t := v.(type) {
			case int:
				ms = t
			case int64:
				ms = int(t)
			case float64:
				ms = int(t)
			case string:
				if n, err := strconv.Atoi(strings.TrimSpace(t)); err == nil {
					ms = n
				}
			}
		}
		if ms < 0 {
			ms = 0
		}
		time.Sleep(time.Duration(ms) * time.Millisecond)
		return 0, fmt.Sprintf("slept %dms", ms), "", nil

	default:
		return 1, "", "unsupported action", fmt.Errorf("unsupported action: %s", st.Action)
	}
}

func renderStep(st Step, ctx map[string]string) Step {
	out := st
	out.ID = renderString(out.ID, ctx)
	out.Action = renderString(out.Action, ctx)
	if out.With == nil {
		return out
	}
	m := make(map[string]any, len(out.With))
	for k, v := range out.With {
		m[k] = renderAny(v, ctx)
	}
	out.With = m
	return out
}

func renderAny(v any, ctx map[string]string) any {
	switch t := v.(type) {
	case string:
		return renderString(t, ctx)
	case []any:
		arr := make([]any, 0, len(t))
		for _, it := range t {
			arr = append(arr, renderAny(it, ctx))
		}
		return arr
	case map[string]any:
		m := make(map[string]any, len(t))
		for k, v := range t {
			m[k] = renderAny(v, ctx)
		}
		return m
	default:
		return v
	}
}

func renderString(s string, ctx map[string]string) string {
	res := s
	for k, v := range ctx {
		res = strings.ReplaceAll(res, "${"+k+"}", v)
	}
	return res
}

func applyHTTPHeaders(req *http.Request, raw any) bool {
	hasContentType := false
	m, ok := raw.(map[string]any)
	if !ok {
		return hasContentType
	}
	for k, v := range m {
		key := strings.TrimSpace(k)
		if key == "" {
			continue
		}
		val := strings.TrimSpace(fmt.Sprintf("%v", v))
		if val == "" {
			continue
		}
		req.Header.Set(key, val)
		if strings.EqualFold(key, "Content-Type") {
			hasContentType = true
		}
	}
	return hasContentType
}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	fi, err := in.Stat()
	if err != nil {
		return err
	}
	if !fi.Mode().IsRegular() {
		return fmt.Errorf("source is not regular file")
	}
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	if _, err := io.Copy(out, in); err != nil {
		return err
	}
	return out.Close()
}

func extractZip(src, dst string) error {
	zr, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer zr.Close()
	for _, f := range zr.File {
		clean := filepath.Clean(f.Name)
		if strings.HasPrefix(clean, "..") || filepath.IsAbs(clean) {
			return fmt.Errorf("invalid zip entry path: %s", f.Name)
		}
		target := filepath.Join(dst, clean)
		if f.FileInfo().IsDir() {
			if err := os.MkdirAll(target, 0o755); err != nil {
				return err
			}
			continue
		}
		if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
			return err
		}
		rc, err := f.Open()
		if err != nil {
			return err
		}
		out, err := os.Create(target)
		if err != nil {
			rc.Close()
			return err
		}
		_, cpErr := io.Copy(out, rc)
		rc.Close()
		if err := out.Close(); err != nil {
			return err
		}
		if cpErr != nil {
			return cpErr
		}
	}
	return nil
}

func extractTar(src, dst string) error {
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()
	return untar(f, dst)
}

func extractTarGz(src, dst string) error {
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()
	gz, err := gzip.NewReader(f)
	if err != nil {
		return err
	}
	defer gz.Close()
	return untar(gz, dst)
}

func untar(r io.Reader, dst string) error {
	tr := tar.NewReader(r)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		clean := filepath.Clean(hdr.Name)
		if strings.HasPrefix(clean, "..") || filepath.IsAbs(clean) {
			return fmt.Errorf("invalid tar entry path: %s", hdr.Name)
		}
		target := filepath.Join(dst, clean)
		switch hdr.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(target, 0o755); err != nil {
				return err
			}
		case tar.TypeReg, tar.TypeRegA:
			if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
				return err
			}
			out, err := os.Create(target)
			if err != nil {
				return err
			}
			if _, err := io.Copy(out, tr); err != nil {
				out.Close()
				return err
			}
			if err := out.Close(); err != nil {
				return err
			}
		}
	}
	return nil
}

func parseRequiredPaths(v any) []string {
	res := []string{}
	switch t := v.(type) {
	case []any:
		for _, it := range t {
			res = append(res, strings.TrimSpace(fmt.Sprintf("%v", it)))
		}
	case []string:
		for _, it := range t {
			res = append(res, strings.TrimSpace(it))
		}
	}
	out := make([]string, 0, len(res))
	for _, p := range res {
		if p != "" {
			out = append(out, p)
		}
	}
	return out
}

func lookupPath(root any, path string) (any, bool) {
	parts := strings.Split(path, ".")
	cur := root
	for _, part := range parts {
		part = strings.TrimSpace(part)
		if part == "" {
			return nil, false
		}
		m, ok := cur.(map[string]any)
		if !ok {
			return nil, false
		}
		next, exists := m[part]
		if !exists {
			return nil, false
		}
		cur = next
	}
	return cur, true
}

func parseExpectSuccess(v any) bool {
	switch t := v.(type) {
	case bool:
		return t
	case string:
		return strings.EqualFold(strings.TrimSpace(t), "true")
	default:
		return false
	}
}

func truthyValue(v any) bool {
	switch t := v.(type) {
	case bool:
		return t
	case string:
		return strings.EqualFold(strings.TrimSpace(t), "true")
	default:
		return false
	}
}

// summarizeJSONErrors 从 Cloudflare 风格的 {"success":false,"errors":[...]} 中提取可读错误摘要
func summarizeJSONErrors(payload any) string {
	msgs := []string{}
	if m, ok := payload.(map[string]any); ok {
		if errs, ok := m["errors"].([]any); ok {
			for _, e := range errs {
				switch ev := e.(type) {
				case map[string]any:
					code := fmt.Sprintf("%v", ev["code"])
					msg := fmt.Sprintf("%v", ev["message"])
					msgs = append(msgs, code+" "+msg)
				default:
					msgs = append(msgs, fmt.Sprintf("%v", ev))
				}
			}
		}
	}
	if len(msgs) == 0 {
		return "unknown error"
	}
	return strings.Join(msgs, "; ")
}

func (e *Executor) finishJob(job *models.OpsJob, status, summary string) {
	job.Status = status
	job.Summary = summary
	job.EndedAt = time.Now()
	_ = e.db.Save(job).Error
}

func isJobCancelled(db *gorm.DB, jobID uint) bool {
	var j models.OpsJob
	if err := db.Select("status").First(&j, jobID).Error; err != nil {
		return false
	}
	return strings.EqualFold(strings.TrimSpace(j.Status), "cancelled")
}

func tail(s string, max int) string {
	s = strings.TrimSpace(s)
	if len(s) <= max {
		return s
	}
	return s[len(s)-max:]
}