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

package host

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

// Manager provides controlled access to the host machine.
// It wraps a Sandbox for command execution and adds file system
// operations with path allow-list enforcement.
type Manager struct {
	sandbox    *Sandbox
	allowedDirs []string
}

// NewManager creates a new host Manager.
func NewManager(sandbox *Sandbox) *Manager {
	m := &Manager{sandbox: sandbox}
	if sandbox != nil {
		m.allowedDirs = sandbox.cfg.AllowedDirs
	}
	return m
}

// SetAllowedDirs updates the list of directories accessible for file operations.
func (m *Manager) SetAllowedDirs(dirs []string) {
	m.allowedDirs = dirs
	m.sandbox.cfg.AllowedDirs = dirs
}

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

// ReadFile reads the contents of a file within allowed directories.
func (m *Manager) ReadFile(path string, maxBytes int) (string, error) {
	if maxBytes <= 0 {
		maxBytes = 1024 * 1024
	}
	if err := m.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 (m *Manager) 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 := m.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 (m *Manager) ListDir(path string) ([]DirEntry, error) {
	if err := m.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
}

// DirEntry represents a filesystem directory entry.
type DirEntry struct {
	Name    string `json:"name"`
	IsDir   bool   `json:"is_dir"`
	Size    int64  `json:"size"`
	ModTime string `json:"mod_time"`
}

// SystemInfo returns basic system information.
func (m *Manager) 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,
	}

	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 (m *Manager) DiskUsage(path string) (map[string]interface{}, error) {
	if err := m.validatePath(path); err != nil {
		return nil, err
	}
	info, err := os.Stat(path)
	if err != nil {
		return nil, fmt.Errorf("cannot stat path: %w", err)
	}
	result := map[string]interface{}{
		"path":     path,
		"is_dir":   info.IsDir(),
		"size":     info.Size(),
		"mod_time": info.ModTime().Format(time.RFC3339),
	}
	return result, nil
}

func (m *Manager) validatePath(path string) error {
	absPath, err := filepath.Abs(path)
	if err != nil {
		return fmt.Errorf("cannot resolve path: %w", err)
	}
	if len(m.allowedDirs) == 0 {
		return nil
	}
	for _, allowed := range m.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)
}