post inicial
This commit is contained in:
lastdragon
@@ -0,0 +1,276 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// maxCommandTimeout es el máximo que el modelo puede pedir para un comando.
|
||||
const maxCommandTimeout = 2 * time.Hour
|
||||
|
||||
// CommandSpec describe un proceso a ejecutar.
|
||||
type CommandSpec struct {
|
||||
Name string
|
||||
Args []string
|
||||
Dir string // directorio de trabajo; no cambia el del agente
|
||||
Timeout time.Duration // obligatorio (> 0)
|
||||
TmpDir string // dónde guardar salidas extensas
|
||||
Env []string // variables adicionales (p.ej. secretos para shell)
|
||||
}
|
||||
|
||||
// CommandResult es el resultado de RunCommand.
|
||||
type CommandResult struct {
|
||||
ExitCode int // -1 si no llegó a ejecutarse o fue matado
|
||||
Stdout string
|
||||
Stderr string
|
||||
Duration time.Duration
|
||||
TimedOut bool
|
||||
Err error // error de arranque o de espera (no incluye exit != 0)
|
||||
Spilled []string
|
||||
}
|
||||
|
||||
// commandEnv es el entorno de los procesos hijos: se hereda el del agente
|
||||
// y se fuerzan valores no interactivos para que ningún comando se quede
|
||||
// esperando entrada o abra un paginador.
|
||||
func commandEnv() []string {
|
||||
env := os.Environ()
|
||||
return append(env,
|
||||
"DEBIAN_FRONTEND=noninteractive",
|
||||
"APT_LISTCHANGES_FRONTEND=none",
|
||||
"PAGER=cat", "SYSTEMD_PAGER=cat", "GIT_PAGER=cat", "MANPAGER=cat",
|
||||
"SYSTEMD_COLORS=0", "NO_COLOR=1", "TERM=dumb",
|
||||
)
|
||||
}
|
||||
|
||||
// RunCommand ejecuta un programa con argumentos, sin shell intermedio.
|
||||
//
|
||||
// Entradas: spec.Name/Args (sin interpretación de shell: no hay problemas
|
||||
// de quoting), spec.Dir y spec.Timeout.
|
||||
// Salidas: stdout y stderr por separado, código de salida y si hubo timeout.
|
||||
// Efectos: el proceso corre en su propio grupo; al vencer el timeout o
|
||||
// cancelarse ctx se mata el grupo completo. stdin es /dev/null.
|
||||
func RunCommand(ctx context.Context, spec CommandSpec) CommandResult {
|
||||
ctx, cancel := context.WithTimeout(ctx, spec.Timeout)
|
||||
defer cancel()
|
||||
|
||||
cmd := exec.CommandContext(ctx, spec.Name, spec.Args...)
|
||||
cmd.Dir = spec.Dir
|
||||
cmd.Env = append(commandEnv(), spec.Env...)
|
||||
cmd.Stdin = nil
|
||||
setProcessGroup(cmd)
|
||||
// Si un nieto mantiene abiertas las tuberías, no esperar más de esto.
|
||||
cmd.WaitDelay = 5 * time.Second
|
||||
|
||||
stdout := newSpillWriter(spec.TmpDir, "stdout")
|
||||
stderr := newSpillWriter(spec.TmpDir, "stderr")
|
||||
defer stdout.Close()
|
||||
defer stderr.Close()
|
||||
cmd.Stdout, cmd.Stderr = stdout, stderr
|
||||
|
||||
start := time.Now()
|
||||
err := cmd.Run()
|
||||
res := CommandResult{Duration: time.Since(start), ExitCode: -1}
|
||||
res.Stdout, res.Stderr = stdout.Text(), stderr.Text()
|
||||
for _, p := range []string{stdout.SpillPath(), stderr.SpillPath()} {
|
||||
if p != "" {
|
||||
res.Spilled = append(res.Spilled, p)
|
||||
}
|
||||
}
|
||||
|
||||
var exitErr *exec.ExitError
|
||||
switch {
|
||||
case errors.Is(ctx.Err(), context.DeadlineExceeded):
|
||||
res.TimedOut = true
|
||||
res.Err = fmt.Errorf("timed out after %s", spec.Timeout)
|
||||
case err == nil:
|
||||
res.ExitCode = 0
|
||||
case errors.As(err, &exitErr):
|
||||
res.ExitCode = exitErr.ExitCode()
|
||||
if ctx.Err() != nil {
|
||||
res.Err = ctx.Err()
|
||||
}
|
||||
default:
|
||||
res.Err = err
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// format produce el texto para el modelo.
|
||||
func (r CommandResult) format(header string) string {
|
||||
var b strings.Builder
|
||||
b.WriteString(header)
|
||||
fmt.Fprintf(&b, "exit_code: %d\nduration: %s\n", r.ExitCode, r.Duration.Round(time.Millisecond))
|
||||
if r.TimedOut {
|
||||
b.WriteString("timed_out: true (process group killed)\n")
|
||||
}
|
||||
if r.Err != nil && !r.TimedOut {
|
||||
fmt.Fprintf(&b, "error: %v\n", r.Err)
|
||||
}
|
||||
fmt.Fprintf(&b, "--- stdout ---\n%s", r.Stdout)
|
||||
if r.Stdout != "" && !strings.HasSuffix(r.Stdout, "\n") {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
fmt.Fprintf(&b, "--- stderr ---\n%s", r.Stderr)
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// resolveTimeout aplica el valor por defecto de la política y el máximo.
|
||||
func resolveTimeout(env *Env, seconds int) time.Duration {
|
||||
if seconds <= 0 {
|
||||
return env.Policy.DefaultCommandTimeout
|
||||
}
|
||||
return min(time.Duration(seconds)*time.Second, maxCommandTimeout)
|
||||
}
|
||||
|
||||
func quoteArgv(argv []string) string {
|
||||
parts := make([]string, len(argv))
|
||||
for i, a := range argv {
|
||||
if a == "" || strings.ContainsAny(a, " \t\n'\"\\$`;&|<>(){}*?[]#~") {
|
||||
parts[i] = "'" + strings.ReplaceAll(a, "'", `'\''`) + "'"
|
||||
} else {
|
||||
parts[i] = a
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, " ")
|
||||
}
|
||||
|
||||
// ---- exec ----
|
||||
|
||||
type execTool struct{}
|
||||
|
||||
type execArgs struct {
|
||||
Command string `json:"command"`
|
||||
Args []string `json:"args"`
|
||||
Cwd string `json:"cwd"`
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
}
|
||||
|
||||
func (execTool) Name() string { return "exec" }
|
||||
func (execTool) Description() string {
|
||||
return "Run a program directly (no shell): command + args. Preferred over shell when no pipes/redirections are needed. Output: exit_code, stdout, stderr."
|
||||
}
|
||||
func (execTool) Schema() map[string]any {
|
||||
return obj([]string{"command"}, map[string]any{
|
||||
"command": str("program name or path, e.g. systemctl"),
|
||||
"args": strList("arguments, one per element, e.g. [\"status\",\"nginx\"]"),
|
||||
"cwd": str("working directory (default PROJECT_ROOT)"),
|
||||
"timeout_seconds": num("timeout in seconds (default from policy, max 7200)"),
|
||||
})
|
||||
}
|
||||
|
||||
func (execTool) Run(ctx context.Context, env *Env, raw json.RawMessage) Result {
|
||||
var a execArgs
|
||||
if err := decode(raw, &a); err != nil {
|
||||
return fail("exec", err)
|
||||
}
|
||||
if strings.TrimSpace(a.Command) == "" {
|
||||
return fail("exec", errors.New("command is required"))
|
||||
}
|
||||
argv := append([]string{a.Command}, a.Args...)
|
||||
cwd := env.Guard.Resolve(a.Cwd, "")
|
||||
display := quoteArgv(argv)
|
||||
|
||||
cls, err := env.Guard.CheckArgv(argv, cwd)
|
||||
if err != nil {
|
||||
return fail(display, err)
|
||||
}
|
||||
if env.DryRun && !cls.ReadOnly() {
|
||||
return dryRun(fmt.Sprintf("run in %s: %s (risks: %v)", cwd, display, cls.AllRisks()))
|
||||
}
|
||||
if st, err := os.Stat(cwd); err != nil || !st.IsDir() {
|
||||
return fail(display, fmt.Errorf("working directory %s does not exist", cwd))
|
||||
}
|
||||
|
||||
args := make([]string, len(a.Args))
|
||||
for i, arg := range a.Args {
|
||||
v, err := env.Secrets.Expand(arg)
|
||||
if err != nil {
|
||||
return fail(display, err)
|
||||
}
|
||||
args[i] = v
|
||||
}
|
||||
spec := CommandSpec{Name: a.Command, Args: args, Dir: cwd, Timeout: resolveTimeout(env, a.TimeoutSeconds), TmpDir: tmpDir(env.Root)}
|
||||
return runWithRetry(ctx, env, spec, display, cwd, cls.AllRisks())
|
||||
}
|
||||
|
||||
func commandResult(display, cwd string, r CommandResult) Result {
|
||||
header := fmt.Sprintf("command: %s\ncwd: %s\n", display, cwd)
|
||||
summary := fmt.Sprintf("%s → exit=%d", display, r.ExitCode)
|
||||
if r.TimedOut {
|
||||
summary = display + " → timed out"
|
||||
} else if r.Err != nil && r.ExitCode < 0 {
|
||||
summary = fmt.Sprintf("%s → %v", display, r.Err)
|
||||
}
|
||||
return Result{
|
||||
Output: r.format(header),
|
||||
OK: r.ExitCode == 0 && r.Err == nil,
|
||||
Summary: summary,
|
||||
ExitCode: r.ExitCode,
|
||||
Audit: []string{"command", display},
|
||||
}
|
||||
}
|
||||
|
||||
// ---- shell ----
|
||||
|
||||
type shellTool struct{}
|
||||
|
||||
type shellArgs struct {
|
||||
Script string `json:"script"`
|
||||
Cwd string `json:"cwd"`
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
}
|
||||
|
||||
func (shellTool) Name() string { return "shell" }
|
||||
func (shellTool) Description() string {
|
||||
return "Run a shell expression with bash -lc (pipes, redirections, &&, loops). Use exec instead for single programs. Every command in the script is checked against the policy."
|
||||
}
|
||||
func (shellTool) Schema() map[string]any {
|
||||
return obj([]string{"script"}, map[string]any{
|
||||
"script": str("shell script, e.g. \"journalctl -u nginx -n 100 | grep -i error\""),
|
||||
"cwd": str("working directory (default PROJECT_ROOT)"),
|
||||
"timeout_seconds": num("timeout in seconds (default from policy, max 7200)"),
|
||||
})
|
||||
}
|
||||
|
||||
func (shellTool) Run(ctx context.Context, env *Env, raw json.RawMessage) Result {
|
||||
var a shellArgs
|
||||
if err := decode(raw, &a); err != nil {
|
||||
return fail("shell", err)
|
||||
}
|
||||
if strings.TrimSpace(a.Script) == "" {
|
||||
return fail("shell", errors.New("script is required"))
|
||||
}
|
||||
cwd := env.Guard.Resolve(a.Cwd, "")
|
||||
display := "sh: " + clip(strings.ReplaceAll(a.Script, "\n", "; "), 300)
|
||||
|
||||
cls, err := env.Guard.CheckScript(a.Script, cwd)
|
||||
if err != nil {
|
||||
return fail(display, err)
|
||||
}
|
||||
if env.DryRun && !cls.ReadOnly() {
|
||||
return dryRun(fmt.Sprintf("run shell script in %s: %s (risks: %v)", cwd, a.Script, cls.AllRisks()))
|
||||
}
|
||||
if st, err := os.Stat(cwd); err != nil || !st.IsDir() {
|
||||
return fail(display, fmt.Errorf("working directory %s does not exist", cwd))
|
||||
}
|
||||
|
||||
shell, flag := "bash", "-lc"
|
||||
if _, err := exec.LookPath("bash"); err != nil {
|
||||
shell, flag = "sh", "-c"
|
||||
}
|
||||
script, secretEnv, err := env.Secrets.ShellEnv(a.Script)
|
||||
if err != nil {
|
||||
return fail(display, err)
|
||||
}
|
||||
spec := CommandSpec{Name: shell, Args: []string{flag, script}, Dir: cwd, Timeout: resolveTimeout(env, a.TimeoutSeconds), TmpDir: tmpDir(env.Root), Env: secretEnv}
|
||||
res := runWithRetry(ctx, env, spec, display, cwd, cls.AllRisks())
|
||||
res.Output = strings.Replace(res.Output, "command: "+display, "script:\n"+a.Script, 1)
|
||||
res.Audit = []string{"command", a.Script}
|
||||
return res
|
||||
}
|
||||
Reference in New Issue
Block a user