binlog-parser/vendor/b612.me/wincmd/permission.go

package wincmd

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

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

func StartProcessWithSYS(appPath, cmdLine, workDir string, runas bool) error {
	var (
		sessionId          win32api.HANDLE
		userToken          win32api.TOKEN = 0
		envInfo            win32api.HANDLE
		impersonationToken win32api.HANDLE = 0
		startupInfo        win32api.StartupInfo
		processInfo        win32api.ProcessInformation
		sessionInformation win32api.HANDLE              = win32api.HANDLE(0)
		sessionCount       int                          = 0
		sessionList        []*win32api.WTS_SESSION_INFO = make([]*win32api.WTS_SESSION_INFO, 0)
		err                error
	)
	if err := win32api.WTSEnumerateSessions(0, 0, 1, &sessionInformation, &sessionCount); err != nil {
		return err
	}
	structSize := unsafe.Sizeof(win32api.WTS_SESSION_INFO{})
	current := uintptr(sessionInformation)
	for i := 0; i < sessionCount; i++ {
		sessionList = append(sessionList, (*win32api.WTS_SESSION_INFO)(unsafe.Pointer(current)))
		current += structSize
	}
	if sessionId, err = func() (win32api.HANDLE, error) {
		for i := range sessionList {
			if sessionList[i].State == win32api.WTSActive {
				return sessionList[i].SessionID, nil
			}
		}
		if sessionId, err := win32api.WTSGetActiveConsoleSessionId(); sessionId == 0xFFFFFFFF {
			return 0xFFFFFFFF, fmt.Errorf("get current user session token: call native WTSGetActiveConsoleSessionId: %s", err)
		} else {
			return win32api.HANDLE(sessionId), nil
		}
	}(); err != nil {
		return err
	}

	if err := win32api.WTSQueryUserToken(sessionId, &impersonationToken); err != nil {
		return err
	}

	if err := win32api.DuplicateTokenEx(impersonationToken, 0, 0, int(win32api.SecurityImpersonation), win32api.TokenPrimary, &userToken); err != nil {
		return fmt.Errorf("call native DuplicateTokenEx: %s", err)
	}
	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 {
			userToken = admin.LinkedToken
		}
	}
	if err := win32api.CloseHandle(impersonationToken); err != nil {
		return fmt.Errorf("close windows handle used for token duplication: %s", err)
	}

	if err := win32api.CreateEnvironmentBlock(&envInfo, userToken, 0); err != nil {
		return fmt.Errorf("create environment details for process: %s", err)
	}
	creationFlags := win32api.CREATE_UNICODE_ENVIRONMENT | win32api.CREATE_NEW_CONSOLE
	startupInfo.ShowWindow = 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)
	}
	return nil
}

func GetRunningProcess() ([]map[string]string, error) {
	result := []map[string]string{}
	pHandle, err := win32api.CreateToolhelp32Snapshot(0x2, 0x0)
	if err != nil {
		return result, err
	}
	for {
		var proc win32api.PROCESSENTRY32
		proc.DwSize = win32api.Ulong(unsafe.Sizeof(proc))
		if err := win32api.Process32Next(pHandle, &proc); err == nil {
			bytetmp := proc.SzExeFile[0:]
			var sakura []byte
			for _, v := range bytetmp {
				if v == byte(0) {
					break
				}
				sakura = append(sakura, v)
			}
			result = append(result, map[string]string{"name": string(sakura), "pid": strconv.Itoa(int(proc.Th32ProcessID)), "ppid": fmt.Sprint(int(proc.Th32ParentProcessID))})
		} else {
			break
		}
	}
	win32api.CloseHandle(pHandle)
	return result, nil
}

func IsProcessRunningByPID(pid int) (bool, error) {
	pHandle, err := win32api.CreateToolhelp32Snapshot(0x2, 0x0)
	if err != nil {
		return false, err
	}
	for {
		var proc win32api.PROCESSENTRY32
		proc.DwSize = win32api.Ulong(unsafe.Sizeof(proc))
		if err := win32api.Process32Next(pHandle, &proc); err == nil {
			bytetmp := int(proc.Th32ProcessID)
			if bytetmp == pid {
				return true, nil
			}
		} else {
			break
		}
	}
	win32api.CloseHandle(pHandle)
	return false, err
}
func IsProcessRunning(name string) (bool, error) {
	pHandle, err := win32api.CreateToolhelp32Snapshot(0x2, 0x0)
	if err != nil {
		return false, err
	}
	for {
		var proc win32api.PROCESSENTRY32
		proc.DwSize = win32api.Ulong(unsafe.Sizeof(proc))
		if err := win32api.Process32Next(pHandle, &proc); err == nil {
			bytetmp := proc.SzExeFile[0:]
			var sakura []byte
			for _, v := range bytetmp {
				if v == byte(0) {
					break
				}
				sakura = append(sakura, v)
			}
			if strings.ToLower(strings.TrimSpace(string(sakura))) == strings.ToLower(strings.TrimSpace(name)) {
				return true, nil
			}
		} else {
			break
		}
	}
	win32api.CloseHandle(pHandle)
	return false, err
}

func GetProcessCount(name string) (int, error) {
	var res int = 0
	pHandle, err := win32api.CreateToolhelp32Snapshot(0x2, 0x0)
	if err != nil {
		return 0, err
	}
	for {
		var proc win32api.PROCESSENTRY32
		proc.DwSize = win32api.Ulong(unsafe.Sizeof(proc))
		if err := win32api.Process32Next(pHandle, &proc); err == nil {
			bytetmp := proc.SzExeFile[0:]
			var sakura []byte
			for _, v := range bytetmp {
				if v == byte(0) {
					break
				}
				sakura = append(sakura, v)
			}
			if strings.ToLower(strings.TrimSpace(string(sakura))) == strings.ToLower(strings.TrimSpace(name)) {
				res++
			}
		} else {
			break
		}
	}
	win32api.CloseHandle(pHandle)
	return res, err
}

func Isas() bool {
	_, errs := registry.OpenKey(registry.LOCAL_MACHINE, `SYSTEM`, registry.ALL_ACCESS)
	if errs != nil {
		return false
	}
	return true
}

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 = 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
	}
	return int(win32api.GetProcessId(sakura.HProcess)), 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
	}
	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
	}
	if _, i, _ := reg.GetStringValue(key); i == 0 {
		return true, nil
	}
	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
	}
	if sa, _, _ := reg.GetStringValue(key); sa == path {
		return true, err
	}
	return false, err
}