wincmd/permission.go

package wincmd

import (
	"fmt"
	"strconv"
	"strings"
	"syscall"
	"unsafe"

	"b612.me/win32api"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/registry"
)

func getActiveSessionID() (win32api.DWORD, error) {
	sessionID, err := win32api.ActiveSessionID()
	if err != nil {
		return 0, fmt.Errorf("resolve active session id: %w", err)
	}
	if sessionID == win32api.WTS_CURRENT_SESSION {
		return 0, fmt.Errorf("active session id is invalid: %#x", sessionID)
	}
	return sessionID, nil
}

func destroyEnvironmentBlock(env win32api.HANDLE) error {
	proc, err := syscall.LoadDLL("userenv.dll")
	if err != nil {
		return err
	}
	defer proc.Release()
	destroy, err := proc.FindProc("DestroyEnvironmentBlock")
	if err != nil {
		return err
	}
	r, _, errno := syscall.Syscall(destroy.Addr(), 1, uintptr(env), 0, 0)
	if r == 0 {
		if errno != 0 {
			return error(errno)
		}
		return syscall.EINVAL
	}
	return nil
}

func StartProcessWithSYS(appPath, cmdLine, workDir string, runas bool) error {
	var (
		sessionId          win32api.DWORD
		userToken          win32api.TOKEN
		envInfo            win32api.HANDLE
		impersonationToken win32api.HANDLE
		startupInfo        win32api.StartupInfo
		processInfo        win32api.ProcessInformation
	)
	sessionId, err := getActiveSessionID()
	if err != nil {
		return fmt.Errorf("get active session id: %w", err)
	}

	if err := win32api.WTSQueryUserToken(sessionId, &impersonationToken); err != nil {
		return err
	}
	defer func() {
		if impersonationToken != 0 {
			_ = win32api.CloseHandle(impersonationToken)
		}
	}()

	if err := win32api.DuplicateTokenEx(impersonationToken, 0, 0, int(win32api.SecurityImpersonation), win32api.TokenPrimary, &userToken); err != nil {
		return fmt.Errorf("call native DuplicateTokenEx: %s", err)
	}
	defer func() {
		if userToken != 0 {
			_ = win32api.CloseHandle(win32api.HANDLE(userToken))
		}
	}()
	if runas {
		var admin win32api.TOKEN_LINKED_TOKEN
		var dt uintptr = 0
		if err := win32api.GetTokenInformation(impersonationToken, 19, uintptr(unsafe.Pointer(&admin)), uintptr(unsafe.Sizeof(admin)), &dt); err == nil && admin.LinkedToken != 0 {
			if userToken != 0 && userToken != admin.LinkedToken {
				_ = win32api.CloseHandle(win32api.HANDLE(userToken))
			}
			userToken = admin.LinkedToken
		}
	}

	if err := win32api.CreateEnvironmentBlock(&envInfo, userToken, 0); err != nil {
		return fmt.Errorf("create environment details for process: %s", err)
	}
	defer func() {
		if envInfo != 0 {
			_ = destroyEnvironmentBlock(envInfo)
		}
	}()
	creationFlags := win32api.CREATE_UNICODE_ENVIRONMENT | win32api.CREATE_NEW_CONSOLE
	startupInfo.Cb = uint32(unsafe.Sizeof(startupInfo))
	startupInfo.ShowWindow = uint16(win32api.SW_SHOW)
	startupInfo.Desktop = windows.StringToUTF16Ptr("winsta0\\default")
	if err := win32api.CreateProcessAsUser(userToken, appPath, cmdLine, 0, 0, 0,
		creationFlags, envInfo, workDir, &startupInfo, &processInfo); err != nil {
		return fmt.Errorf("create process as user: %s", err)
	}
	if processInfo.Process != 0 {
		_ = win32api.CloseHandle(processInfo.Process)
	}
	if processInfo.Thread != 0 {
		_ = win32api.CloseHandle(processInfo.Thread)
	}
	return nil
}

func processImageName(proc windows.ProcessEntry32) string {
	return windows.UTF16ToString(proc.ExeFile[:])
}

func walkProcesses(fn func(proc windows.ProcessEntry32) (bool, error)) error {
	if fn == nil {
		return nil
	}
	pHandle, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0)
	if err != nil {
		return err
	}
	defer func() {
		_ = windows.CloseHandle(pHandle)
	}()

	var proc windows.ProcessEntry32
	proc.Size = uint32(unsafe.Sizeof(proc))
	if err := windows.Process32First(pHandle, &proc); err != nil {
		if errno, ok := err.(syscall.Errno); ok && errno == syscall.ERROR_NO_MORE_FILES {
			return nil
		}
		return err
	}

	for {
		stop, err := fn(proc)
		if err != nil {
			return err
		}
		if stop {
			return nil
		}

		if err := windows.Process32Next(pHandle, &proc); err != nil {
			if errno, ok := err.(syscall.Errno); ok && errno == syscall.ERROR_NO_MORE_FILES {
				return nil
			}
			return err
		}
	}
}

func GetRunningProcess() ([]map[string]string, error) {
	result := []map[string]string{}
	err := walkProcesses(func(proc windows.ProcessEntry32) (bool, error) {
		result = append(result, map[string]string{
			"name": processImageName(proc),
			"pid":  strconv.Itoa(int(proc.ProcessID)),
			"ppid": fmt.Sprint(int(proc.ParentProcessID)),
		})
		return false, nil
	})
	return result, err
}

func IsProcessRunningByPID(pid int) (bool, error) {
	found := false
	err := walkProcesses(func(proc windows.ProcessEntry32) (bool, error) {
		if int(proc.ProcessID) == pid {
			found = true
			return true, nil
		}
		return false, nil
	})
	return found, err
}
func IsProcessRunning(name string) (bool, error) {
	target := strings.TrimSpace(name)
	found := false
	err := walkProcesses(func(proc windows.ProcessEntry32) (bool, error) {
		if strings.EqualFold(strings.TrimSpace(processImageName(proc)), target) {
			found = true
			return true, nil
		}
		return false, nil
	})
	return found, err
}

func GetProcessCount(name string) (int, error) {
	var count int
	target := strings.TrimSpace(name)
	err := walkProcesses(func(proc windows.ProcessEntry32) (bool, error) {
		if strings.EqualFold(strings.TrimSpace(processImageName(proc)), target) {
			count++
		}
		return false, nil
	})
	return count, err
}

// IsElevated reports whether the current process token is elevated and belongs to local Administrators.
func IsElevated() (bool, error) {
	var token windows.Token
	if err := windows.OpenProcessToken(windows.CurrentProcess(), windows.TOKEN_QUERY, &token); err != nil {
		return false, err
	}
	defer token.Close()

	elevated := token.IsElevated()
	inAdminGroup, err := isCurrentUserInAdminGroup(token)
	if err != nil {
		if elevated {
			return true, nil
		}
		return false, err
	}
	return elevated && inAdminGroup, nil
}

func isCurrentUserInAdminGroup(token windows.Token) (bool, error) {
	adminSID, err := windows.CreateWellKnownSid(windows.WinBuiltinAdministratorsSid)
	if err != nil {
		return false, err
	}
	member, err := token.IsMember(adminSID)
	if err == nil {
		return member, nil
	}
	// CheckTokenMembership supports Token(0) fallback to caller's effective token.
	return windows.Token(0).IsMember(adminSID)
}

func Isas() bool {
	elevated, err := IsElevated()
	if err != nil {
		return false
	}
	return elevated
}

func StartProcess(appPath, cmdLine, wordDir string, runas bool, ShowWindow int) error {
	var cst string
	if runas {
		cst = "runas"
	} else {
		cst = "open"
	}
	r := win32api.ShellExecute(0, cst, appPath, cmdLine, wordDir, ShowWindow)
	return r
}

func StartProcessWithPID(appPath, cmdLine, workDir string, runas bool, ShowWindow int) (int, error) {
	var sakura win32api.SHELLEXECUTEINFOW
	sakura.Hwnd = 0
	sakura.NShow = int32(ShowWindow)
	sakura.FMask = 0x00000040
	sakura.LpParameters = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(cmdLine)))
	sakura.LpFile = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(appPath)))
	sakura.LpDirectory = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(workDir)))
	if runas {
		sakura.LpVerb = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("runas")))
	} else {
		sakura.LpVerb = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("open")))
	}
	sakura.CbSize = win32api.DWORD(unsafe.Sizeof(sakura))

	if err := win32api.ShellExecuteEx(&sakura); err != nil {
		return 0, err
	}
	pid := int(win32api.GetProcessId(sakura.HProcess))
	if sakura.HProcess != 0 {
		_ = win32api.CloseHandle(sakura.HProcess)
	}
	return pid, nil
}

func AutoRun(key, path string) (bool, error) {
	reg, errs := registry.OpenKey(registry.LOCAL_MACHINE, `Software\Microsoft\Windows\CurrentVersion\Run`, registry.ALL_ACCESS)
	if errs != nil {
		return false, errs
	}
	defer reg.Close()
	if errs = reg.SetStringValue(key, path); errs != nil {
		return false, errs
	}
	return true, nil
}

func DeleteAutoRun(key string) (bool, error) {
	reg, errs := registry.OpenKey(registry.LOCAL_MACHINE, `Software\Microsoft\Windows\CurrentVersion\Run`, registry.ALL_ACCESS)
	if errs != nil {
		return false, errs
	}
	defer reg.Close()
	if _, _, err := reg.GetStringValue(key); err != nil {
		if err == registry.ErrNotExist {
			return true, nil
		}
		return false, err
	}
	if errs = reg.DeleteValue(key); errs != nil {
		return false, errs
	}
	return true, nil
}

func IsAutoRun(key, path string) (bool, error) {
	reg, err := registry.OpenKey(registry.LOCAL_MACHINE, `Software\Microsoft\Windows\CurrentVersion\Run`, registry.ALL_ACCESS)
	if err != nil {
		return false, err
	}
	defer reg.Close()
	sa, _, err := reg.GetStringValue(key)
	if err != nil {
		if err == registry.ErrNotExist {
			return false, nil
		}
		return false, err
	}
	return sa == path, nil
}