Tries to replicate qwinsta but return structured objects

Get-WTSSessionInfo.ps1

# Copyright: (c) 2022, Jordan Borean (@jborean93) <[email protected]>
# MIT License (see LICENSE or https://opensource.org/licenses/MIT)

<#Check if loaded to make dot-source testing easier#>
if(-not ('Wtsapi32.Native' -as [type])){
    Add-Type -TypeDefinition @'
using System;
using System.Runtime.InteropServices;

namespace Wtsapi32
{
    public enum WtsConnectState
    {
        Active,
        Connected,
        ConnectQuery,
        Shadow,
        Disconnected,
        Idle,
        Listen,
        Reset,
        Down,
        Init
    }

    public class SessionInfo
    {
        public int SessionId { get; set; }
        public WtsConnectState State { get; set; }
        public string SessionName { get; set; }
        public string HostName { get; set; }
        public string UserName { get; set; }
        public string DomainName { get; set; }
        public string FarmName { get; set; }
        public string ComputerName { get; set; }
    }

    public class Native
    {
        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        private struct WTS_SESSION_INFO_1W
        {
            public int ExecEnvId;
            public WtsConnectState State;
            public int SessionId;
            public string pSessionName;
            public string pHostName;
            public string pUserName;
            public string pDomainName;
            public string pFarmName;
        }

        [DllImport("Wtsapi32.dll")]
        public static extern void WTSCloseServer(IntPtr hServer);

        [DllImport("Wtsapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
        private static extern bool WTSEnumerateSessionsExW(
            IntPtr hServer,
            ref int pLevel,
            int Filter,
            out IntPtr ppSessionInfo,
            out int pCount);

        public static SessionInfo[] WTSEnumerateSessionsEx(System.Collections.Hashtable serverInfo)
        {
            IntPtr serverHandle = (IntPtr)serverInfo["serverHandle"];
            string serverName = (string)serverInfo["serverName"];
            int level = 1;
            IntPtr rawInfo;
            int count;
            if (!WTSEnumerateSessionsExW(serverHandle, ref level, 0, out rawInfo, out count))
                throw new System.ComponentModel.Win32Exception();
            try
            {
                var infoCollection = new System.Collections.Generic.List<SessionInfo>();
                IntPtr currentOffset = rawInfo;
                for (int i = 0; i < count; i++)
                {
                    WTS_SESSION_INFO_1W info = Marshal.PtrToStructure<WTS_SESSION_INFO_1W>(currentOffset);
                    infoCollection.Add(new SessionInfo()
                    {
                        SessionId = info.SessionId,
                        State = info.State,
                        SessionName = info.pSessionName,
                        HostName = info.pHostName,
                        UserName = info.pUserName,
                        DomainName = info.pDomainName,
                        FarmName = info.pFarmName,
                        ComputerName = serverName
                    });
                    currentOffset = IntPtr.Add(currentOffset, Marshal.SizeOf(typeof(WTS_SESSION_INFO_1W)));
                }
                return infoCollection.ToArray();
            }
            finally
            {
                WTSFreeMemoryExW(2, rawInfo, count);
            }
        }

        [DllImport("Wtsapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
        private static extern bool WTSFreeMemoryExW(
            int WTSTypeClass,
            IntPtr pMemory,
            int NumberOfEntries);

        [DllImport("Wtsapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
        private static extern IntPtr WTSOpenServerExW(string pServerName);

        public static System.Collections.Hashtable WTSOpenServerEx(string serverName)
        {
            IntPtr serverHandle = WTSOpenServerExW(serverName);
            if (serverHandle == IntPtr.Zero)
                throw new System.ComponentModel.Win32Exception();
            return new System.Collections.Hashtable() { { "serverHandle", serverHandle }, { "serverName", serverName } };
        }

        [DllImport("Wtsapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
        public static extern bool WTSLogoffSession(IntPtr hServer, int SessionId);
    }
}
'@
}

Function Get-WTSSession {
    <#
    .SYNOPSIS
    Enumerates sessions on a Windows host.

    .DESCRIPTION
    Enumerates all the sessions available on a Windows host through the WTSEnumerateSessionsExW API.

    .PARAMETER ComputerName
    A list of hosts to query the sessions for. Omit or set to an empty array to check the local host.

    .EXAMPLE
    Get-WTSSessionInfo

    .NOTES
    The output object is modeled after https://docs.microsoft.com/en-us/windows/win32/api/wtsapi32/ns-wtsapi32-wts_session_info_1w.
    #>
    [Alias('gwts')]
    [CmdletBinding()]
    param(
        [Parameter(ValueFromPipeline, ValueFromPipelineByPropertyName)]
        [string[]] $ComputerName = $env:COMPUTERNAME
    )
    process {
        foreach ($name in $ComputerName) {
            try {
                $serverInfo = [Wtsapi32.Native]::WTSOpenServerEx($name)
                [Wtsapi32.Native]::WTSEnumerateSessionsEx($serverInfo)
            }
            catch {
                Write-Error -ErrorRecord $_
            }
            finally {
                [Wtsapi32.Native]::WTSCloseServer($serverInfo.serverHandle)
            }
        }
    }
}
function Remove-WTSSession {
    <#
    .SYNOPSIS
    Removes a session from a Windows host.

    .DESCRIPTION
    Removes a session from a Windows host through the WTSLogoffSession API.

    .PARAMETER SessionInfo
    The session info object returned from Get-WTSSessionInfo.

    .EXAMPLE
    Get-WTSSessionInfo | Where-Object { $_.State -eq 'Disconnected' } | Remove-WTSSession

    .EXAMPLE
    Get-WTSSessionInfo | Remove-WTSSession

    #>
    [Alias('rwts')]
    [CmdletBinding(SupportsShouldProcess,ConfirmImpact = 'High')]
    param(
        [Parameter(ValueFromPipeline)]
        [Wtsapi32.SessionInfo] $SessionInfo
    )
    process {
        if (-not $PSCmdlet.ShouldProcess("Logoff Session $($SessionInfo.SessionId) for $($SessionInfo.UserName)@$($SessionInfo.ComputerName) State: $($SessionInfo.State)")) {
            return
        }
        $serverInfo = [Wtsapi32.Native]::WTSOpenServerEx($SessionInfo.ComputerName)
        try {
            [void][Wtsapi32.Native]::WTSLogoffSession($serverInfo.serverHandle, $SessionInfo.SessionId)
        }
        catch {
            Write-Error -ErrorRecord $_
        }
        finally {
            [Wtsapi32.Native]::WTSCloseServer($serverInfo.serverHandle)
        }
    }
}