AetherForge/agent/deploy/desktop_path.go

package deploy

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

const desktopPathPrefix = "@desktop/"

// UserDesktopDir returns the interactive user's Desktop folder for the current OS.
func UserDesktopDir() (string, error) {
	switch runtime.GOOS {
	case "windows":
		return windowsDesktopDir()
	case "darwin":
		return unixDesktopFromHome("Desktop")
	default:
		return linuxDesktopDir()
	}
}

func windowsDesktopDir() (string, error) {
	profile := strings.TrimSpace(os.Getenv("USERPROFILE"))
	if profile == "" {
		return "", fmt.Errorf("USERPROFILE not set")
	}
	candidates := []string{
		filepath.Join(profile, "Desktop"),
		filepath.Join(profile, "OneDrive", "Desktop"),
		filepath.Join(profile, "OneDrive - Personal", "Desktop"),
	}
	for _, c := range candidates {
		if st, err := os.Stat(c); err == nil && st.IsDir() {
			return filepath.Clean(c), nil
		}
	}
	// Create default Desktop if missing (unusual profiles)
	fallback := candidates[0]
	if err := os.MkdirAll(fallback, 0o755); err != nil {
		return "", err
	}
	return fallback, nil
}

func unixDesktopFromHome(sub string) (string, error) {
	home, err := os.UserHomeDir()
	if err != nil || home == "" {
		return "", fmt.Errorf("home directory unavailable")
	}
	desktop := filepath.Join(home, sub)
	if st, err := os.Stat(desktop); err == nil && st.IsDir() {
		return filepath.Clean(desktop), nil
	}
	if err := os.MkdirAll(desktop, 0o755); err != nil {
		return "", fmt.Errorf("desktop: %w", err)
	}
	return desktop, nil
}

func linuxDesktopDir() (string, error) {
	if out, err := exec.Command("xdg-user-dir", "DESKTOP").Output(); err == nil {
		p := strings.TrimSpace(string(out))
		if p != "" {
			if st, err := os.Stat(p); err == nil && st.IsDir() {
				return filepath.Clean(p), nil
			}
		}
	}
	return unixDesktopFromHome("Desktop")
}

// ResolveDesktopFile joins a sanitized filename onto the user Desktop.
func ResolveDesktopFile(filename string) (string, error) {
	desktop, err := UserDesktopDir()
	if err != nil {
		return "", err
	}
	name := sanitizeDesktopFilename(filename)
	if name == "" {
		name = "upload.bin"
	}
	return filepath.Join(desktop, name), nil
}

func sanitizeDesktopFilename(name string) string {
	name = strings.TrimSpace(name)
	name = strings.ReplaceAll(name, "\\", "/")
	if name == "" {
		return ""
	}
	// Allow subfolders under Desktop but block traversal.
	parts := strings.Split(name, "/")
	var clean []string
	for _, p := range parts {
		p = strings.TrimSpace(p)
		if p == "" || p == "." || p == ".." {
			continue
		}
		clean = append(clean, p)
	}
	return filepath.Join(clean...)
}

func remotePathHasTraversal(remote string) bool {
	remote = strings.TrimSpace(remote)
	if remote == "" {
		return false
	}
	if strings.HasPrefix(remote, "~/") {
		remote = remote[2:]
	}
	remote = strings.ReplaceAll(remote, "\\", "/")
	for _, part := range strings.Split(remote, "/") {
		if part == ".." {
			return true
		}
	}
	return false
}

// ResolveRemotePath expands @desktop/…, desktop:…, and ~/… for upload/download commands.
func ResolveRemotePath(remote string) (string, error) {
	remote = strings.TrimSpace(remote)
	if remote == "" {
		return "", fmt.Errorf("remote path is empty")
	}
	if remotePathHasTraversal(remote) {
		return "", fmt.Errorf("path traversal (..) is not allowed")
	}
	lower := strings.ToLower(remote)
	if strings.HasPrefix(lower, "desktop:") {
		return ResolveDesktopFile(remote[len("desktop:"):])
	}
	if strings.HasPrefix(remote, desktopPathPrefix) {
		return ResolveDesktopFile(remote[len(desktopPathPrefix):])
	}
	if strings.HasPrefix(remote, "~/") {
		home, err := os.UserHomeDir()
		if err != nil {
			return "", err
		}
		return filepath.Clean(filepath.Join(home, remote[2:])), nil
	}
	return filepath.Clean(remote), nil
}