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

package host

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

// DirectBackend executes commands directly on the host via os/exec,
// with command allowlist and directory restrictions for safety.
type DirectBackend struct {
	sandbox     *Sandbox
	allowedDirs []string
}

// NewDirectBackend creates a host execution backend that runs commands
// directly on the host machine with sandbox restrictions.
func NewDirectBackend(sandbox *Sandbox) *DirectBackend {
	b := &DirectBackend{sandbox: sandbox}
	if sandbox != nil {
		b.allowedDirs = sandbox.cfg.AllowedDirs
	}
	return b
}

func (b *DirectBackend) Name() string { return "direct" }

// SetAllowedDirs updates the directories accessible for file operations.
func (b *DirectBackend) SetAllowedDirs(dirs []string) {
	b.allowedDirs = dirs
	if b.sandbox != nil {
		b.sandbox.cfg.AllowedDirs = dirs
	}
}

// Exec runs a command in the sandbox.
func (b *DirectBackend) Exec(ctx context.Context, command, workDir string, timeout time.Duration) (*ExecResult, error) {
	return b.sandbox.Exec(ctx, command, workDir, timeout)
}

// ReadFile reads the contents of a file within allowed directories.
func (b *DirectBackend) ReadFile(path string, maxBytes int) (string, error) {
	if maxBytes <= 0 {
		maxBytes = 1024 * 1024
	}
	if err := b.validatePath(path); err != nil {
		return "", err
	}
	info, err := os.Stat(path)
	if err != nil {
		return "", fmt.Errorf("cannot stat file: %w", err)
	}
	if info.IsDir() {
		return "", fmt.Errorf("path is a directory: %s", path)
	}
	if info.Size() > int64(maxBytes) {
		return "", fmt.Errorf("file too large: %d bytes (max %d)", info.Size(), maxBytes)
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("cannot read file: %w", err)
	}
	if len(data) > maxBytes {
		data = data[:maxBytes]
	}
	return string(data), nil
}

// WriteFile writes data to a file within allowed directories.
func (b *DirectBackend) WriteFile(path, content string, maxBytes int) error {
	if maxBytes <= 0 {
		maxBytes = 1024 * 1024
	}
	if len(content) > maxBytes {
		return fmt.Errorf("content too large: %d bytes (max %d)", len(content), maxBytes)
	}
	if err := b.validatePath(path); err != nil {
		return err
	}
	dir := filepath.Dir(path)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return fmt.Errorf("cannot create directory: %w", err)
	}
	return os.WriteFile(path, []byte(content), 0644)
}

// ListDir lists directory contents within allowed directories.
func (b *DirectBackend) ListDir(path string) ([]DirEntry, error) {
	if err := b.validatePath(path); err != nil {
		return nil, err
	}
	entries, err := os.ReadDir(path)
	if err != nil {
		return nil, fmt.Errorf("cannot read directory: %w", err)
	}
	result := make([]DirEntry, 0, len(entries))
	for _, e := range entries {
		info, _ := e.Info()
		size := int64(0)
		modTime := time.Time{}
		if info != nil {
			size = info.Size()
			modTime = info.ModTime()
		}
		result = append(result, DirEntry{
			Name:    e.Name(),
			IsDir:   e.IsDir(),
			Size:    size,
			ModTime: modTime.Format(time.RFC3339),
		})
	}
	return result, nil
}

// SystemInfo returns basic system information.
func (b *DirectBackend) SystemInfo() map[string]interface{} {
	hostname, _ := os.Hostname()
	wd, _ := os.Getwd()

	info := map[string]interface{}{
		"hostname":   hostname,
		"os":         runtime.GOOS,
		"arch":       runtime.GOARCH,
		"num_cpu":    runtime.NumCPU(),
		"go_version": runtime.Version(),
		"work_dir":   wd,
		"backend":    "direct",
	}

	if runtime.GOOS == "windows" {
		cmd := exec.Command("systeminfo")
		out, err := cmd.Output()
		if err == nil {
			lines := strings.Split(string(out), "\n")
			for _, line := range lines {
				line = strings.TrimSpace(line)
				if strings.Contains(line, "Total Physical Memory") {
					parts := strings.SplitN(line, ":", 2)
					if len(parts) == 2 {
						info["total_memory"] = strings.TrimSpace(parts[1])
					}
				}
				if strings.Contains(line, "OS Name") {
					parts := strings.SplitN(line, ":", 2)
					if len(parts) == 2 {
						info["os_name"] = strings.TrimSpace(parts[1])
					}
				}
			}
		}
	} else {
		if data, err := os.ReadFile("/proc/meminfo"); err == nil {
			for _, line := range strings.Split(string(data), "\n") {
				if strings.HasPrefix(line, "MemTotal:") {
					info["total_memory"] = strings.TrimSpace(strings.TrimPrefix(line, "MemTotal:"))
					break
				}
			}
		}
	}
	return info
}

// DiskUsage returns disk usage for the given path.
func (b *DirectBackend) DiskUsage(path string) (map[string]interface{}, error) {
	if err := b.validatePath(path); err != nil {
		return nil, err
	}
	info, err := os.Stat(path)
	if err != nil {
		return nil, fmt.Errorf("cannot stat path: %w", err)
	}
	return map[string]interface{}{
		"path":     path,
		"is_dir":   info.IsDir(),
		"size":     info.Size(),
		"mod_time": info.ModTime().Format(time.RFC3339),
	}, nil
}

func (b *DirectBackend) validatePath(path string) error {
	absPath, err := filepath.Abs(path)
	if err != nil {
		return fmt.Errorf("cannot resolve path: %w", err)
	}
	if len(b.allowedDirs) == 0 {
		return nil
	}
	for _, allowed := range b.allowedDirs {
		absAllowed, err := filepath.Abs(allowed)
		if err != nil {
			continue
		}
		if strings.HasPrefix(absPath, absAllowed+string(os.PathSeparator)) || absPath == absAllowed {
			return nil
		}
	}
	return fmt.Errorf("path not in allowed directories: %s", path)
}