Cyrene/backend/ai-core/internal/host/sandbox.go

package host

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

// SandboxConfig configures the sandbox execution environment.
type SandboxConfig struct {
	AllowedCommands []string
	AllowedDirs     []string
	MaxOutputBytes  int
	DefaultTimeout  time.Duration
	MaxTimeout      time.Duration
}

// DefaultSandboxConfig returns a safe default configuration.
func DefaultSandboxConfig() SandboxConfig {
	return SandboxConfig{
		AllowedCommands: []string{
			"echo", "cat", "ls", "dir", "pwd", "date", "time",
			"wc", "head", "tail", "sort", "uniq", "grep", "find",
			"python", "python3", "node", "go", "rustc", "cargo",
			"git", "curl", "wget", "ping", "nslookup", "tracert",
			"dotnet", "java", "javac", "gcc", "g++", "make", "cmake",
			"npm", "npx", "yarn", "pnpm", "pip", "pip3",
			"docker", "kubectl", "helm",
			"ffmpeg", "ffprobe", "imagemagick", "convert",
			"systeminfo", "tasklist", "taskkill", "netstat",
		},
		MaxOutputBytes: 512 * 1024,
		DefaultTimeout: 30 * time.Second,
		MaxTimeout:     300 * time.Second,
	}
}

// Sandbox provides a safe execution environment for host commands.
type Sandbox struct {
	cfg SandboxConfig
}

// NewSandbox creates a new sandbox.
func NewSandbox(cfg SandboxConfig) *Sandbox {
	return &Sandbox{cfg: cfg}
}

// ExecResult holds the result of a sandboxed command execution.
type ExecResult struct {
	Stdout   string `json:"stdout"`
	Stderr   string `json:"stderr"`
	ExitCode int    `json:"exit_code"`
	Duration string `json:"duration"`
	TimedOut bool   `json:"timed_out"`
}

// Exec runs a command inside the sandbox. The command string is parsed into
// the executable name and arguments. Returns the combined output.
func (s *Sandbox) Exec(ctx context.Context, command string, workDir string, timeout time.Duration) (*ExecResult, error) {
	if command == "" {
		return nil, fmt.Errorf("empty command")
	}

	parts := strings.Fields(command)
	if len(parts) == 0 {
		return nil, fmt.Errorf("empty command")
	}

	cmdName := parts[0]
	var args []string
	if len(parts) > 1 {
		args = parts[1:]
	}

	if !s.isCommandAllowed(cmdName) {
		return nil, fmt.Errorf("command not allowed: %s", cmdName)
	}

	if workDir == "" {
		workDir = s.defaultWorkDir()
	}
	if err := s.validateWorkDir(workDir); err != nil {
		return nil, err
	}

	if timeout <= 0 {
		timeout = s.cfg.DefaultTimeout
	}
	if timeout > s.cfg.MaxTimeout {
		timeout = s.cfg.MaxTimeout
	}

	execCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	cmd := exec.CommandContext(execCtx, cmdName, args...)
	cmd.Dir = workDir
	cmd.Env = s.filteredEnv()

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	start := time.Now()
	err := cmd.Run()
	elapsed := time.Since(start)

	result := &ExecResult{
		Duration: elapsed.Round(time.Millisecond).String(),
	}

	if stdout.Len() > s.cfg.MaxOutputBytes {
		result.Stdout = stdout.String()[:s.cfg.MaxOutputBytes] + "\n... [output truncated]"
	} else {
		result.Stdout = stdout.String()
	}
	if stderr.Len() > s.cfg.MaxOutputBytes {
		result.Stderr = stderr.String()[:s.cfg.MaxOutputBytes] + "\n... [output truncated]"
	} else {
		result.Stderr = stderr.String()
	}

	if execCtx.Err() == context.DeadlineExceeded {
		result.TimedOut = true
		result.ExitCode = -1
		return result, fmt.Errorf("command timed out after %s", timeout)
	}

	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			result.ExitCode = exitErr.ExitCode()
		} else {
			result.ExitCode = -1
		}
	} else {
		result.ExitCode = 0
	}

	return result, err
}

func (s *Sandbox) isCommandAllowed(cmd string) bool {
	if len(s.cfg.AllowedCommands) == 0 {
		return true
	}
	base := filepath.Base(cmd)
	base = strings.TrimSuffix(base, ".exe")
	for _, allowed := range s.cfg.AllowedCommands {
		if strings.EqualFold(base, allowed) {
			return true
		}
	}
	return false
}

func (s *Sandbox) validateWorkDir(dir string) error {
	info, err := os.Stat(dir)
	if err != nil {
		return fmt.Errorf("work directory not accessible: %s: %w", dir, err)
	}
	if !info.IsDir() {
		return fmt.Errorf("not a directory: %s", dir)
	}

	if len(s.cfg.AllowedDirs) == 0 {
		return nil
	}

	absDir, err := filepath.Abs(dir)
	if err != nil {
		return fmt.Errorf("cannot resolve path: %w", err)
	}

	for _, allowed := range s.cfg.AllowedDirs {
		absAllowed, err := filepath.Abs(allowed)
		if err != nil {
			continue
		}
		if strings.HasPrefix(absDir, absAllowed) {
			return nil
		}
	}
	return fmt.Errorf("directory not in allowed list: %s", dir)
}

func (s *Sandbox) defaultWorkDir() string {
	if len(s.cfg.AllowedDirs) > 0 {
		return s.cfg.AllowedDirs[0]
	}
	wd, _ := os.Getwd()
	return wd
}

func (s *Sandbox) filteredEnv() []string {
	allowed := map[string]bool{
		"PATH": true, "HOME": true, "USER": true, "USERNAME": true,
		"TMP": true, "TEMP": true, "TMPDIR": true,
		"LANG": true, "LC_ALL": true, "SHELL": true,
		"SYSTEMROOT": true, "WINDIR": true, "ProgramFiles": true,
		"GOPATH": true, "GOROOT": true, "GOPROXY": true,
		"NODE_PATH": true, "PYTHONPATH": true,
		"JAVA_HOME": true, "DOTNET_ROOT": true,
		"CARGO_HOME": true, "RUSTUP_HOME": true,
	}
	var filtered []string
	for _, e := range os.Environ() {
		kv := strings.SplitN(e, "=", 2)
		if len(kv) == 2 && allowed[kv[0]] {
			filtered = append(filtered, e)
		}
	}
	filtered = append(filtered, "CYRENE_SANDBOX=1")
	return filtered
}