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

package host

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

// WSLBackend executes commands inside a WSL2 distribution,
// providing a full Linux OS environment isolated from the Windows host.
type WSLBackend struct {
	distro   string
	username string
	password string
	timeout  time.Duration

	userEnsured bool
}

// NewWSLBackend creates a WSL backend that runs commands in the
// specified WSL distribution as the given user. On first use,
// the user is automatically created with sudo privileges.
func NewWSLBackend(distro, username, password string, defaultTimeout time.Duration) *WSLBackend {
	if defaultTimeout <= 0 {
		defaultTimeout = 30 * time.Second
	}
	if username == "" {
		username = "cyrene"
	}
	return &WSLBackend{
		distro:   distro,
		username: username,
		password: password,
		timeout:  defaultTimeout,
	}
}

func (b *WSLBackend) Name() string { return "wsl" }

// ensureUser creates the configured user inside the WSL distro on first call.
// The user gets sudo privileges and the configured password.
func (b *WSLBackend) ensureUser() error {
	if b.userEnsured {
		return nil
	}

	// Check if user already exists
	checkCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	checkCmd := exec.CommandContext(checkCtx, "wsl.exe", "-d", b.distro, "--", "id", b.username)
	if checkCmd.Run() == nil {
		b.userEnsured = true
		return nil
	}

	// Create user with home directory, set password, add to sudo group
	// If password is empty, create user without password (sudo won't need it
	// if NOPASSWD is configured, but we still set a random one for safety)
	pwd := b.password
	if pwd == "" {
		pwd = "cyrene"
	}

	// Escape single quotes in password for the shell echo command
	escapedPwd := strings.ReplaceAll(pwd, "'", "'\\''")
	script := fmt.Sprintf(
		"useradd -m -s /bin/bash %s && echo '%s:%s' | chpasswd && usermod -aG sudo %s",
		b.username, b.username, escapedPwd, b.username,
	)

	createCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	createCmd := exec.CommandContext(createCtx, "wsl.exe", "-d", b.distro, "--", "bash", "-c", script)
	if out, err := createCmd.CombinedOutput(); err != nil {
		return fmt.Errorf("cannot create user %s: %s — %w", b.username, string(out), err)
	}

	b.userEnsured = true
	return nil
}

// Exec runs a command inside the WSL distribution via bash.
func (b *WSLBackend) Exec(ctx context.Context, command, workDir string, timeout time.Duration) (*ExecResult, error) {
	if command == "" {
		return nil, fmt.Errorf("empty command")
	}

	if err := b.ensureUser(); err != nil {
		return nil, err
	}

	if timeout <= 0 {
		timeout = b.timeout
	}

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

	// Build the bash command to run inside WSL
	script := command
	if workDir != "" {
		wslPath := windowsToWSLPath(workDir)
		script = fmt.Sprintf("cd %s && %s", shellEscape(wslPath), command)
	}

	cmd := exec.CommandContext(execCtx, "wsl.exe", "-d", b.distro, "--", "bash", "-c", script)

	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(),
		Stdout:   stdout.String(),
		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
}

// ReadFile reads a file from the WSL filesystem using cat.
func (b *WSLBackend) ReadFile(path string, maxBytes int) (string, error) {
	if maxBytes <= 0 {
		maxBytes = 1024 * 1024
	}
	if err := b.ensureUser(); err != nil {
		return "", err
	}
	wslPath := windowsToWSLPath(path)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	cmd := exec.CommandContext(ctx, "wsl.exe", "-d", b.distro, "--", "cat", wslPath)
	out, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("cannot read file %s: %w", path, err)
	}
	if len(out) > maxBytes {
		out = out[:maxBytes]
	}
	return string(out), nil
}

// WriteFile writes content to a file in the WSL filesystem.
func (b *WSLBackend) 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.ensureUser(); err != nil {
		return err
	}
	wslPath := windowsToWSLPath(path)
	// Create parent directory first
	dir := filepath.Dir(wslPath)
	_ = exec.Command("wsl.exe", "-d", b.distro, "--", "mkdir", "-p", dir).Run()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	cmd := exec.CommandContext(ctx, "wsl.exe", "-d", b.distro, "--", "bash", "-c",
		fmt.Sprintf("cat > %s", shellEscape(wslPath)))
	cmd.Stdin = strings.NewReader(content)
	_, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("cannot write file %s: %w", path, err)
	}
	return nil
}

// ListDir lists a directory in the WSL filesystem using ls.
func (b *WSLBackend) ListDir(path string) ([]DirEntry, error) {
	if err := b.ensureUser(); err != nil {
		return nil, err
	}
	wslPath := windowsToWSLPath(path)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	cmd := exec.CommandContext(ctx, "wsl.exe", "-d", b.distro, "--", "bash", "-c",
		fmt.Sprintf("stat -c '%%n|%%F|%%s|%%Y' %s/* 2>/dev/null || echo ''", shellEscape(wslPath)))
	out, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("cannot list dir %s: %w", path, err)
	}

	lines := strings.Split(strings.TrimSpace(string(out)), "\n")
	result := make([]DirEntry, 0, len(lines))
	for _, line := range lines {
		if line == "" {
			continue
		}
		parts := strings.SplitN(line, "|", 4)
		if len(parts) < 4 {
			continue
		}
		var size int64
		fmt.Sscanf(parts[2], "%d", &size)
		var modTimeUnix int64
		fmt.Sscanf(parts[3], "%d", &modTimeUnix)
		modTime := time.Unix(modTimeUnix, 0).Format(time.RFC3339)
		isDir := strings.Contains(parts[1], "directory")
		result = append(result, DirEntry{
			Name:    filepath.Base(parts[0]),
			IsDir:   isDir,
			Size:    size,
			ModTime: modTime,
		})
	}
	return result, nil
}

// SystemInfo returns system information from inside the WSL distribution.
func (b *WSLBackend) SystemInfo() map[string]interface{} {
	info := map[string]interface{}{
		"backend": "wsl",
		"distro":  b.distro,
	}

	if err := b.ensureUser(); err != nil {
		info["error"] = err.Error()
		return info
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// uname
	if out, err := exec.CommandContext(ctx, "wsl.exe", "-d", b.distro, "--", "uname", "-a").Output(); err == nil {
		info["uname"] = strings.TrimSpace(string(out))
	}

	// hostname
	if out, err := exec.CommandContext(ctx, "wsl.exe", "-d", b.distro, "--", "hostname").Output(); err == nil {
		info["hostname"] = strings.TrimSpace(string(out))
	}

	// memory info
	if out, err := exec.CommandContext(ctx, "wsl.exe", "-d", b.distro, "--", "free", "-h").Output(); err == nil {
		info["memory"] = strings.TrimSpace(string(out))
	}

	// disk info
	if out, err := exec.CommandContext(ctx, "wsl.exe", "-d", b.distro, "--", "df", "-h", "/").Output(); err == nil {
		lines := strings.Split(strings.TrimSpace(string(out)), "\n")
		if len(lines) > 1 {
			info["disk"] = strings.TrimSpace(lines[1])
		}
	}

	return info
}

// DiskUsage returns disk usage for a path inside WSL.
func (b *WSLBackend) DiskUsage(path string) (map[string]interface{}, error) {
	wslPath := windowsToWSLPath(path)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	cmd := exec.CommandContext(ctx, "wsl.exe", "-d", b.distro, "--", "stat", wslPath)
	out, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("cannot stat path %s: %w", path, err)
	}

	// Parse stat output minimally
	result := map[string]interface{}{
		"path":   path,
		"wsl_path": wslPath,
		"stat":   strings.TrimSpace(string(out)),
	}
	return result, nil
}

// windowsToWSLPath converts a Windows path to its WSL equivalent.
// C:\Users\foo → /mnt/c/Users/foo
// If the path is already a WSL path (starts with /), return as-is.
func windowsToWSLPath(path string) string {
	if strings.HasPrefix(path, "/") {
		return path // Already a Unix path
	}
	// Handle Windows drive letter: C:\... → /mnt/c/...
	if len(path) >= 2 && path[1] == ':' {
		drive := strings.ToLower(string(path[0]))
		rest := strings.TrimPrefix(path[2:], "\\")
		rest = strings.ReplaceAll(rest, "\\", "/")
		return fmt.Sprintf("/mnt/%s/%s", drive, rest)
	}
	return path
}

// shellEscape escapes a string for safe use in a shell command.
func shellEscape(s string) string {
	// Use single quotes and escape any single quotes in the string
	escaped := strings.ReplaceAll(s, "'", "'\\''")
	return "'" + escaped + "'"
}