Add terminal service session info (#1525)

pkg/headers/wtsapi32/wtsapi32.go

@@ -0,0 +1,198 @@
+package wtsapi32
+
+import (
+	"fmt"
+	"syscall"
+	"unsafe"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"golang.org/x/sys/windows"
+)
+
+type WTSTypeClass int
+
+// The valid values for the WTSTypeClass enumeration
+const (
+	WTSTypeProcessInfoLevel0 WTSTypeClass = iota
+	WTSTypeProcessInfoLevel1
+	WTSTypeSessionInfoLevel1
+)
+
+type WTSConnectState uint32
+
+const (
+	// wtsActive A user is logged on to the WinStation. This state occurs when a user is signed in and actively connected to the device.
+	wtsActive WTSConnectState = iota
+	// wtsConnected The WinStation is connected to the client.
+	wtsConnected
+	// wtsConnectQuery The WinStation is in the process of connecting to the client.
+	wtsConnectQuery
+	// wtsShadow The WinStation is shadowing another WinStation.
+	wtsShadow
+	// wtsDisconnected The WinStation is active but the client is disconnected.
+	// This state occurs when a user is signed in but not actively connected to the device, such as when the user has chosen to exit to the lock screen.
+	wtsDisconnected
+	// wtsIdle The WinStation is waiting for a client to connect.
+	wtsIdle
+	// wtsListen The WinStation is listening for a connection. A listener session waits for requests for new client connections.
+	// No user is logged on a listener session. A listener session cannot be reset, shadowed, or changed to a regular client session.
+	wtsListen
+	// wtsReset The WinStation is being reset.
+	wtsReset
+	// wtsDown The WinStation is down due to an error.
+	wtsDown
+	// wtsInit The WinStation is initializing.
+	wtsInit
+)
+
+// WTSSessionInfo1w contains information about a session on a Remote Desktop Session Host (RD Session Host) server.
+// docs: https://docs.microsoft.com/en-us/windows/win32/api/wtsapi32/ns-wtsapi32-wts_session_info_1w
+type wtsSessionInfo1 struct {
+	// ExecEnvID An identifier that uniquely identifies the session within the list of sessions returned by the WTSEnumerateSessionsEx function.
+	ExecEnvID uint32
+	// State A value of the WTSConnectState enumeration type that specifies the connection state of a Remote Desktop Services session.
+	State uint32
+	// SessionID A session identifier assigned by the RD Session Host server, RD Virtualization Host server, or virtual machine.
+	SessionID uint32
+	// pSessionName A pointer to a null-terminated string that contains the name of this session. For example, "services", "console", or "RDP-Tcp#0".
+	pSessionName *uint16
+	// pHostName A pointer to a null-terminated string that contains the name of the computer that the session is running on.
+	// If the session is running directly on an RD Session Host server or RD Virtualization Host server, the string contains NULL.
+	// If the session is running on a virtual machine, the string contains the name of the virtual machine.
+	pHostName *uint16
+	// pUserName A pointer to a null-terminated string that contains the name of the user who is logged on to the session.
+	// If no user is logged on to the session, the string contains NULL.
+	pUserName *uint16
+	// pDomainName A pointer to a null-terminated string that contains the domain name of the user who is logged on to the session.
+	// If no user is logged on to the session, the string contains NULL.
+	pDomainName *uint16
+	// pFarmName A pointer to a null-terminated string that contains the name of the farm that the virtual machine is joined to.
+	// If the session is not running on a virtual machine that is joined to a farm, the string contains NULL.
+	pFarmName *uint16
+}
+
+type WTSSession struct {
+	ExecEnvID   uint32
+	State       WTSConnectState
+	SessionID   uint32
+	SessionName string
+	HostName    string
+	UserName    string
+	DomainName  string
+	FarmName    string
+}
+
+var (
+	wtsapi32 = windows.NewLazySystemDLL("wtsapi32.dll")
+
+	procWTSOpenServerEx        = wtsapi32.NewProc("WTSOpenServerExW")
+	procWTSEnumerateSessionsEx = wtsapi32.NewProc("WTSEnumerateSessionsExW")
+	procWTSFreeMemoryEx        = wtsapi32.NewProc("WTSFreeMemoryExW")
+	procWTSCloseServer         = wtsapi32.NewProc("WTSCloseServer")
+
+	WTSSessionStates = map[WTSConnectState]string{
+		wtsActive:       "active",
+		wtsConnected:    "connected",
+		wtsConnectQuery: "connect_query",
+		wtsShadow:       "shadow",
+		wtsDisconnected: "disconnected",
+		wtsIdle:         "idle",
+		wtsListen:       "listen",
+		wtsReset:        "reset",
+		wtsDown:         "down",
+		wtsInit:         "init",
+	}
+)
+
+func WTSOpenServer(server string) (syscall.Handle, error) {
+	var (
+		err        error
+		serverName *uint16
+	)
+
+	if server != "" {
+		serverName, err = syscall.UTF16PtrFromString(server)
+		if err != nil {
+			return syscall.InvalidHandle, err
+		}
+	}
+
+	r1, _, err := procWTSOpenServerEx.Call(uintptr(unsafe.Pointer(serverName)))
+	serverHandle := syscall.Handle(r1)
+
+	if serverHandle == syscall.InvalidHandle {
+		return syscall.InvalidHandle, err
+	}
+
+	return serverHandle, nil
+}
+
+func WTSCloseServer(server syscall.Handle) error {
+	_, _, err := procWTSCloseServer.Call(uintptr(server))
+	if err != nil {
+		return fmt.Errorf("failed to close server: %w", err)
+	}
+
+	return err
+}
+
+func WTSFreeMemoryEx(class WTSTypeClass, pMemory uintptr, NumberOfEntries uint32) error {
+	_, _, err := procWTSFreeMemoryEx.Call(
+		uintptr(class),
+		pMemory,
+		uintptr(NumberOfEntries),
+	)
+
+	return err
+}
+
+func WTSEnumerateSessionsEx(server syscall.Handle, logger log.Logger) ([]WTSSession, error) {
+	var sessionInfoPointer uintptr
+	var count uint32
+
+	pLevel := uint32(1)
+	r1, _, err := procWTSEnumerateSessionsEx.Call(
+		uintptr(server),
+		uintptr(unsafe.Pointer(&pLevel)),
+		uintptr(0),
+		uintptr(unsafe.Pointer(&sessionInfoPointer)),
+		uintptr(unsafe.Pointer(&count)),
+	)
+
+	if r1 != 1 {
+		return nil, err
+	}
+
+	if sessionInfoPointer != 0 {
+		defer func(class WTSTypeClass, pMemory uintptr, NumberOfEntries uint32) {
+			err := WTSFreeMemoryEx(class, pMemory, NumberOfEntries)
+			if err != nil {
+				_ = level.Error(logger).Log("msg", "failed to free memory", "err", err)
+			}
+		}(WTSTypeSessionInfoLevel1, sessionInfoPointer, count)
+	}
+
+	var sizeTest wtsSessionInfo1
+	sessionSize := unsafe.Sizeof(sizeTest)
+
+	sessions := make([]WTSSession, 0, count)
+	for i := uint32(0); i < count; i++ {
+		curPtr := unsafe.Pointer(sessionInfoPointer + (uintptr(i) * sessionSize))
+		data := (*wtsSessionInfo1)(curPtr)
+
+		sessionInfo := WTSSession{
+			ExecEnvID:   data.ExecEnvID,
+			State:       WTSConnectState(data.State),
+			SessionID:   data.SessionID,
+			SessionName: windows.UTF16PtrToString(data.pSessionName),
+			HostName:    windows.UTF16PtrToString(data.pHostName),
+			UserName:    windows.UTF16PtrToString(data.pUserName),
+			DomainName:  windows.UTF16PtrToString(data.pDomainName),
+			FarmName:    windows.UTF16PtrToString(data.pFarmName),
+		}
+		sessions = append(sessions, sessionInfo)
+	}
+
+	return sessions, nil
+}