437 lines
13 KiB
Go
437 lines
13 KiB
Go
//go:build windows
|
|
|
|
package windowsx
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"strings"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
"unsafe"
|
|
|
|
"github.com/example/sessionguard/internal/model"
|
|
"golang.org/x/sys/windows"
|
|
"golang.org/x/sys/windows/registry"
|
|
)
|
|
|
|
var (
|
|
wtsapi32 = windows.NewLazySystemDLL("wtsapi32.dll")
|
|
procWTSEnumerateSessionsW = wtsapi32.NewProc("WTSEnumerateSessionsW")
|
|
procWTSEnumerateProcessesW = wtsapi32.NewProc("WTSEnumerateProcessesW")
|
|
procWTSFreeMemory = wtsapi32.NewProc("WTSFreeMemory")
|
|
procWTSQuerySessionInformationW = wtsapi32.NewProc("WTSQuerySessionInformationW")
|
|
procWTSQueryUserToken = wtsapi32.NewProc("WTSQueryUserToken")
|
|
procWTSLogoffSession = wtsapi32.NewProc("WTSLogoffSession")
|
|
procWTSDisconnectSession = wtsapi32.NewProc("WTSDisconnectSession")
|
|
procWTSSendMessageW = wtsapi32.NewProc("WTSSendMessageW")
|
|
userenv = windows.NewLazySystemDLL("userenv.dll")
|
|
procDeleteProfileW = userenv.NewProc("DeleteProfileW")
|
|
kernel32 = windows.NewLazySystemDLL("kernel32.dll")
|
|
procGetTickCount64 = kernel32.NewProc("GetTickCount64")
|
|
procGlobalMemoryStatusEx = kernel32.NewProc("GlobalMemoryStatusEx")
|
|
procExpandEnvironmentStringsW = kernel32.NewProc("ExpandEnvironmentStringsW")
|
|
procGetSystemTimes = kernel32.NewProc("GetSystemTimes")
|
|
procGetDiskFreeSpaceExW = kernel32.NewProc("GetDiskFreeSpaceExW")
|
|
psapi = windows.NewLazySystemDLL("psapi.dll")
|
|
procGetProcessMemoryInfo = psapi.NewProc("GetProcessMemoryInfo")
|
|
)
|
|
|
|
type wtsSessionInfo struct {
|
|
SessionID uint32
|
|
WinStationName *uint16
|
|
State uint32
|
|
}
|
|
type wtsProcessInfo struct {
|
|
SessionID uint32
|
|
ProcessID uint32
|
|
ProcessName *uint16
|
|
UserSID *windows.SID
|
|
}
|
|
type wtsClientAddress struct {
|
|
AddressFamily uint32
|
|
Address [20]byte
|
|
}
|
|
|
|
type wtsInfoExLevel1 struct {
|
|
SessionID uint32
|
|
SessionState uint32
|
|
SessionFlags int32
|
|
WinStationName [33]uint16
|
|
UserName [21]uint16
|
|
DomainName [18]uint16
|
|
LogonTime int64
|
|
ConnectTime int64
|
|
DisconnectTime int64
|
|
LastInputTime int64
|
|
CurrentTime int64
|
|
IncomingBytes uint32
|
|
OutgoingBytes uint32
|
|
IncomingFrames uint32
|
|
OutgoingFrames uint32
|
|
}
|
|
|
|
const (
|
|
wtsUserName = 5
|
|
wtsWinStationName = 6
|
|
wtsDomainName = 7
|
|
wtsClientName = 10
|
|
wtsClientAddressClass = 14
|
|
wtsSessionInfoEx = 25
|
|
)
|
|
|
|
var stateNames = map[uint32]string{0: "Active", 1: "Connected", 2: "ConnectQuery", 3: "Shadow", 4: "Disconnected", 5: "Idle", 6: "Listen", 7: "Reset", 8: "Down", 9: "Init"}
|
|
|
|
func Sessions() ([]model.Session, error) {
|
|
var buf uintptr
|
|
var count uint32
|
|
r1, _, e := procWTSEnumerateSessionsW.Call(0, 0, 1, uintptr(unsafe.Pointer(&buf)), uintptr(unsafe.Pointer(&count)))
|
|
if r1 == 0 {
|
|
return nil, fmt.Errorf("WTSEnumerateSessionsW: %w", e)
|
|
}
|
|
defer procWTSFreeMemory.Call(buf)
|
|
rows := unsafe.Slice((*wtsSessionInfo)(unsafe.Pointer(buf)), int(count))
|
|
out := make([]model.Session, 0, len(rows))
|
|
for _, row := range rows {
|
|
s := model.Session{ID: row.SessionID, State: stateNames[row.State]}
|
|
if s.State == "" {
|
|
s.State = fmt.Sprintf("State%d", row.State)
|
|
}
|
|
if row.WinStationName != nil {
|
|
s.StationName = windows.UTF16PtrToString(row.WinStationName)
|
|
}
|
|
s.User, _ = queryString(row.SessionID, wtsUserName)
|
|
s.Domain, _ = queryString(row.SessionID, wtsDomainName)
|
|
s.ClientName, _ = queryString(row.SessionID, wtsClientName)
|
|
if s.StationName == "" {
|
|
s.StationName, _ = queryString(row.SessionID, wtsWinStationName)
|
|
}
|
|
if ca, err := queryClientAddress(row.SessionID); err == nil {
|
|
s.ClientAddress = ca
|
|
}
|
|
if ex, err := querySessionInfoEx(row.SessionID); err == nil {
|
|
s.LogonAt = fileTimePtr(ex.LogonTime)
|
|
s.ConnectAt = fileTimePtr(ex.ConnectTime)
|
|
s.LastInputAt = fileTimePtr(ex.LastInputTime)
|
|
if s.State == "Disconnected" {
|
|
s.DisconnectedSince = fileTimePtr(ex.DisconnectTime)
|
|
}
|
|
if s.LastInputAt != nil {
|
|
idle := time.Since(*s.LastInputAt)
|
|
if idle > 0 {
|
|
s.IdleSeconds = int64(idle / time.Second)
|
|
}
|
|
}
|
|
}
|
|
if s.User != "" {
|
|
var token windows.Token
|
|
r, _, _ := procWTSQueryUserToken.Call(uintptr(row.SessionID), uintptr(unsafe.Pointer(&token)))
|
|
if r != 0 {
|
|
if tu, err := token.GetTokenUser(); err == nil && tu.User.Sid != nil {
|
|
s.SID = tu.User.Sid.String()
|
|
}
|
|
_ = token.Close()
|
|
}
|
|
}
|
|
out = append(out, s)
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func queryString(sessionID uint32, class uintptr) (string, error) {
|
|
var p uintptr
|
|
var bytes uint32
|
|
r1, _, e := procWTSQuerySessionInformationW.Call(0, uintptr(sessionID), class, uintptr(unsafe.Pointer(&p)), uintptr(unsafe.Pointer(&bytes)))
|
|
if r1 == 0 {
|
|
return "", e
|
|
}
|
|
defer procWTSFreeMemory.Call(p)
|
|
if p == 0 || bytes < 2 {
|
|
return "", nil
|
|
}
|
|
return windows.UTF16PtrToString((*uint16)(unsafe.Pointer(p))), nil
|
|
}
|
|
func queryClientAddress(sessionID uint32) (string, error) {
|
|
var p uintptr
|
|
var bytes uint32
|
|
r, _, e := procWTSQuerySessionInformationW.Call(0, uintptr(sessionID), wtsClientAddressClass, uintptr(unsafe.Pointer(&p)), uintptr(unsafe.Pointer(&bytes)))
|
|
if r == 0 {
|
|
return "", e
|
|
}
|
|
defer procWTSFreeMemory.Call(p)
|
|
if p == 0 || bytes < uint32(unsafe.Sizeof(wtsClientAddress{})) {
|
|
return "", nil
|
|
}
|
|
a := (*wtsClientAddress)(unsafe.Pointer(p))
|
|
if a.AddressFamily == 2 {
|
|
return fmt.Sprintf("%d.%d.%d.%d", a.Address[2], a.Address[3], a.Address[4], a.Address[5]), nil
|
|
}
|
|
return "", nil
|
|
}
|
|
func querySessionInfoEx(sessionID uint32) (wtsInfoExLevel1, error) {
|
|
var out wtsInfoExLevel1
|
|
var p uintptr
|
|
var bytes uint32
|
|
r, _, e := procWTSQuerySessionInformationW.Call(0, uintptr(sessionID), wtsSessionInfoEx, uintptr(unsafe.Pointer(&p)), uintptr(unsafe.Pointer(&bytes)))
|
|
if r == 0 {
|
|
return out, e
|
|
}
|
|
defer procWTSFreeMemory.Call(p)
|
|
if p == 0 || bytes < 8+uint32(unsafe.Sizeof(out)) {
|
|
return out, fmt.Errorf("WTSSessionInfoEx buffer too small")
|
|
}
|
|
level := *(*uint32)(unsafe.Pointer(p))
|
|
if level != 1 {
|
|
return out, fmt.Errorf("unsupported WTSSessionInfoEx level %d", level)
|
|
}
|
|
out = *(*wtsInfoExLevel1)(unsafe.Pointer(p + 8))
|
|
return out, nil
|
|
}
|
|
func fileTimePtr(v int64) *time.Time {
|
|
if v <= 0 {
|
|
return nil
|
|
}
|
|
const unixDelta = 116444736000000000
|
|
ns := (v - unixDelta) * 100
|
|
if ns <= 0 {
|
|
return nil
|
|
}
|
|
t := time.Unix(0, ns).UTC()
|
|
return &t
|
|
}
|
|
|
|
type memoryStatusEx struct {
|
|
Length uint32
|
|
MemoryLoad uint32
|
|
TotalPhys uint64
|
|
AvailPhys uint64
|
|
TotalPageFile uint64
|
|
AvailPageFile uint64
|
|
TotalVirtual uint64
|
|
AvailVirtual uint64
|
|
AvailExtendedVirtual uint64
|
|
}
|
|
type filetime struct {
|
|
LowDateTime uint32
|
|
HighDateTime uint32
|
|
}
|
|
|
|
var cpuMu sync.Mutex
|
|
var prevIdle, prevKernel, prevUser uint64
|
|
|
|
func ft64(f filetime) uint64 { return uint64(f.HighDateTime)<<32 | uint64(f.LowDateTime) }
|
|
func cpuPercent() float64 {
|
|
var idle, kernel, user filetime
|
|
r, _, _ := procGetSystemTimes.Call(uintptr(unsafe.Pointer(&idle)), uintptr(unsafe.Pointer(&kernel)), uintptr(unsafe.Pointer(&user)))
|
|
if r == 0 {
|
|
return 0
|
|
}
|
|
i, k, u := ft64(idle), ft64(kernel), ft64(user)
|
|
cpuMu.Lock()
|
|
defer cpuMu.Unlock()
|
|
pi, pk, pu := prevIdle, prevKernel, prevUser
|
|
prevIdle, prevKernel, prevUser = i, k, u
|
|
if pk == 0 {
|
|
return 0
|
|
}
|
|
total := (k - pk) + (u - pu)
|
|
if total == 0 {
|
|
return 0
|
|
}
|
|
busy := total - (i - pi)
|
|
return float64(busy) * 100 / float64(total)
|
|
}
|
|
|
|
func Server() (model.ServerInfo, error) {
|
|
host, _ := os.Hostname()
|
|
m := memoryStatusEx{Length: uint32(unsafe.Sizeof(memoryStatusEx{}))}
|
|
r, _, e := procGlobalMemoryStatusEx.Call(uintptr(unsafe.Pointer(&m)))
|
|
if r == 0 {
|
|
return model.ServerInfo{}, fmt.Errorf("GlobalMemoryStatusEx: %w", e)
|
|
}
|
|
ticks, _, _ := procGetTickCount64.Call()
|
|
info := model.ServerInfo{Hostname: host, OS: "Windows", UptimeSeconds: uint64(ticks) / 1000, MemoryTotal: m.TotalPhys, MemoryAvailable: m.AvailPhys, CPUPercent: cpuPercent()}
|
|
drive := os.Getenv("SystemDrive")
|
|
if drive == "" {
|
|
drive = `C:`
|
|
}
|
|
root, err := windows.UTF16PtrFromString(drive + `\`)
|
|
if err == nil {
|
|
var avail, total, free uint64
|
|
if rr, _, _ := procGetDiskFreeSpaceExW.Call(uintptr(unsafe.Pointer(root)), uintptr(unsafe.Pointer(&avail)), uintptr(unsafe.Pointer(&total)), uintptr(unsafe.Pointer(&free))); rr != 0 {
|
|
info.DiskTotal = total
|
|
info.DiskFree = free
|
|
}
|
|
}
|
|
if k, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\Microsoft\Windows NT\CurrentVersion`, registry.QUERY_VALUE); err == nil {
|
|
defer k.Close()
|
|
if v, _, err := k.GetStringValue("ProductName"); err == nil {
|
|
info.OS = v
|
|
}
|
|
if v, _, err := k.GetStringValue("DisplayVersion"); err == nil {
|
|
info.Version = v
|
|
}
|
|
if v, _, err := k.GetStringValue("CurrentBuildNumber"); err == nil {
|
|
info.Build = v
|
|
}
|
|
}
|
|
return info, nil
|
|
}
|
|
|
|
func Processes() ([]model.ProcessInfo, error) {
|
|
var buf uintptr
|
|
var count uint32
|
|
r, _, e := procWTSEnumerateProcessesW.Call(0, 0, 1, uintptr(unsafe.Pointer(&buf)), uintptr(unsafe.Pointer(&count)))
|
|
if r == 0 {
|
|
return nil, fmt.Errorf("WTSEnumerateProcessesW: %w", e)
|
|
}
|
|
defer procWTSFreeMemory.Call(buf)
|
|
rows := unsafe.Slice((*wtsProcessInfo)(unsafe.Pointer(buf)), int(count))
|
|
out := make([]model.ProcessInfo, 0, len(rows))
|
|
for _, row := range rows {
|
|
p := model.ProcessInfo{PID: row.ProcessID, SessionID: row.SessionID}
|
|
if row.ProcessName != nil {
|
|
p.Name = windows.UTF16PtrToString(row.ProcessName)
|
|
}
|
|
if row.UserSID != nil {
|
|
p.UserSID = row.UserSID.String()
|
|
}
|
|
p.MemoryBytes = processMemory(row.ProcessID)
|
|
out = append(out, p)
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
type processMemoryCounters struct {
|
|
CB uint32
|
|
PageFaultCount uint32
|
|
PeakWorkingSetSize uintptr
|
|
WorkingSetSize uintptr
|
|
QuotaPeakPagedPoolUsage uintptr
|
|
QuotaPagedPoolUsage uintptr
|
|
QuotaPeakNonPagedPoolUsage uintptr
|
|
QuotaNonPagedPoolUsage uintptr
|
|
PagefileUsage uintptr
|
|
PeakPagefileUsage uintptr
|
|
}
|
|
|
|
func processMemory(pid uint32) uint64 {
|
|
h, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION|windows.PROCESS_VM_READ, false, pid)
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
defer windows.CloseHandle(h)
|
|
m := processMemoryCounters{CB: uint32(unsafe.Sizeof(processMemoryCounters{}))}
|
|
r, _, _ := procGetProcessMemoryInfo.Call(uintptr(h), uintptr(unsafe.Pointer(&m)), uintptr(m.CB))
|
|
if r == 0 {
|
|
return 0
|
|
}
|
|
return uint64(m.WorkingSetSize)
|
|
}
|
|
|
|
func ProfilePath(sid string) (string, error) {
|
|
k, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\Microsoft\Windows NT\CurrentVersion\ProfileList\`+sid, registry.QUERY_VALUE)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer k.Close()
|
|
p, _, err := k.GetStringValue("ProfileImagePath")
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return expandEnv(p), nil
|
|
}
|
|
func expandEnv(s string) string {
|
|
in, err := windows.UTF16PtrFromString(s)
|
|
if err != nil {
|
|
return s
|
|
}
|
|
n, _, _ := procExpandEnvironmentStringsW.Call(uintptr(unsafe.Pointer(in)), 0, 0)
|
|
if n == 0 {
|
|
return s
|
|
}
|
|
buf := make([]uint16, n)
|
|
n2, _, _ := procExpandEnvironmentStringsW.Call(uintptr(unsafe.Pointer(in)), uintptr(unsafe.Pointer(&buf[0])), uintptr(n))
|
|
if n2 == 0 || n2 > n {
|
|
return s
|
|
}
|
|
return windows.UTF16ToString(buf)
|
|
}
|
|
func DeleteProfile(sid string) error {
|
|
p, err := windows.UTF16PtrFromString(sid)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
r, _, e := procDeleteProfileW.Call(uintptr(unsafe.Pointer(p)), 0, 0)
|
|
if r == 0 {
|
|
if e == syscall.Errno(0) {
|
|
return fmt.Errorf("DeleteProfileW failed")
|
|
}
|
|
return fmt.Errorf("DeleteProfileW(%s): %w", sid, e)
|
|
}
|
|
return nil
|
|
}
|
|
func LogoffSession(id uint32) error {
|
|
return boolCallSession(procWTSLogoffSession, "WTSLogoffSession", id)
|
|
}
|
|
func DisconnectSession(id uint32) error {
|
|
return boolCallSession(procWTSDisconnectSession, "WTSDisconnectSession", id)
|
|
}
|
|
func boolCallSession(p *windows.LazyProc, name string, id uint32) error {
|
|
r, _, e := p.Call(0, uintptr(id), 0)
|
|
if r == 0 {
|
|
if e == syscall.Errno(0) {
|
|
return fmt.Errorf("%s(%d) failed", name, id)
|
|
}
|
|
return fmt.Errorf("%s(%d): %w", name, id, e)
|
|
}
|
|
return nil
|
|
}
|
|
func SendMessage(sessionID uint32, title, message string) error {
|
|
if strings.TrimSpace(title) == "" {
|
|
title = "SessionGuard"
|
|
}
|
|
t, err := windows.UTF16FromString(title)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
m, err := windows.UTF16FromString(message)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var response uint32
|
|
r, _, e := procWTSSendMessageW.Call(0, uintptr(sessionID), uintptr(unsafe.Pointer(&t[0])), uintptr((len(t)-1)*2), uintptr(unsafe.Pointer(&m[0])), uintptr((len(m)-1)*2), 0, 60, uintptr(unsafe.Pointer(&response)), 0)
|
|
if r == 0 {
|
|
if e == syscall.Errno(0) {
|
|
return fmt.Errorf("WTSSendMessageW(%d) failed", sessionID)
|
|
}
|
|
return fmt.Errorf("WTSSendMessageW(%d): %w", sessionID, e)
|
|
}
|
|
return nil
|
|
}
|
|
func TerminateProcess(pid uint32) error {
|
|
h, err := windows.OpenProcess(windows.PROCESS_TERMINATE, false, pid)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer windows.CloseHandle(h)
|
|
return windows.TerminateProcess(h, 1)
|
|
}
|
|
func RestartServer(reason string) error {
|
|
if strings.TrimSpace(reason) == "" {
|
|
reason = "SessionGuard maintenance restart"
|
|
}
|
|
return exec.Command("shutdown.exe", "/r", "/t", "0", "/d", "p:4:1", "/c", reason).Run()
|
|
}
|
|
func MachineID() (string, error) {
|
|
k, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\Microsoft\Cryptography`, registry.QUERY_VALUE)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer k.Close()
|
|
v, _, err := k.GetStringValue("MachineGuid")
|
|
return strings.TrimSpace(v), err
|
|
}
|