insight/src/Agent/Insight.Agent/Network/Handlers/SessionHandler.cs

using Insight.Domain.Interfaces;
using Insight.Domain.Network;
using Insight.Domain.Network.Agent.Messages;
using System.Runtime.InteropServices;
using System.Runtime.Versioning;

namespace Insight.Agent.Network.Handlers;

[SupportedOSPlatform("windows")]
public class SessionHandler : IMessageHandler<AgentSession>
{
    public async ValueTask HandleAsync<TMessage>(AgentSession sender, TMessage message, CancellationToken cancellationToken) where TMessage : IMessage
    {
        switch (message)
        {
            case InventoryRequest:
                {
                    var result = new Collection<Session>();
                    result.AddRange(GetSessions());

                    await sender.SendAsync(result, cancellationToken);
                    break;
                }
        }
    }

    private static List<Session> GetSessions()
    {
        var sessions = new List<Session>();

        foreach (var s in NativeMethods.GetSessions())
        {
            sessions.Add(new Session
            {
                Sid = s.SessionId.ToString(),
                User = s.Username,
                Type = s.Workstation,
                Status = s.State.ToString(),
                Remote = s.IPAddress
            });
        }

        return sessions;
    }

    private static partial class NativeMethods
    {
        //public const int WTS_CURRENT_SESSION = -1;

        [DllImport("wtsapi32.dll")]
        static extern int WTSEnumerateSessions(
        nint pServer,
        [MarshalAs(UnmanagedType.U4)] int iReserved,
        [MarshalAs(UnmanagedType.U4)] int iVersion,
        ref nint pSessionInfo,
        [MarshalAs(UnmanagedType.U4)] ref int iCount);

        [DllImport("Wtsapi32.dll")]
        private static extern bool WTSQuerySessionInformation(
            nint pServer,
            int iSessionID,
            WTS_INFO_CLASS oInfoClass,
            out nint pBuffer,
            out uint iBytesReturned);

        [DllImport("wtsapi32.dll")]
        static extern void WTSFreeMemory(
            nint pMemory);

        [StructLayout(LayoutKind.Sequential)]
        private struct WTS_CLIENT_ADDRESS
        {
            public int iAddressFamily;
            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 20)]
            public byte[] bAddress;
        }

        [StructLayout(LayoutKind.Sequential)]
        private struct WTS_SESSION_INFO
        {
            public int iSessionID;
            [MarshalAs(UnmanagedType.LPStr)]
            public string sWinsWorkstationName;
            public WTS_CONNECTSTATE_CLASS oState;
        }

        [StructLayout(LayoutKind.Sequential)]
        private struct WTS_CLIENT_DISPLAY
        {
            public int iHorizontalResolution;
            public int iVerticalResolution;
            //1 = The display uses 4 bits per pixel for a maximum of 16 colors.
            //2 = The display uses 8 bits per pixel for a maximum of 256 colors.
            //4 = The display uses 16 bits per pixel for a maximum of 2^16 colors.
            //8 = The display uses 3-byte RGB values for a maximum of 2^24 colors.
            //16 = The display uses 15 bits per pixel for a maximum of 2^15 colors.
            public int iColorDepth;
        }

        public enum WTS_CONNECTSTATE_CLASS
        {
            WTSActive,
            WTSConnected,
            WTSConnectQuery,
            WTSShadow,
            WTSDisconnected,
            WTSIdle,
            WTSListen,
            WTSReset,
            WTSDown,
            WTSInit
        }

        public enum WTS_INFO_CLASS
        {
            WTSInitialProgram,
            WTSApplicationName,
            WTSWorkingDirectory,
            WTSOEMId,
            WTSSessionId,
            WTSUserName,
            WTSWinStationName,
            WTSDomainName,
            WTSConnectState,
            WTSClientBuildNumber,
            WTSClientName,
            WTSClientDirectory,
            WTSClientProductId,
            WTSClientHardwareId,
            WTSClientAddress,
            WTSClientDisplay,
            WTSClientProtocolType,
            WTSIdleTime,
            WTSLogonTime,
            WTSIncomingBytes,
            WTSOutgoingBytes,
            WTSIncomingFrames,
            WTSOutgoingFrames,
            WTSClientInfo,
            WTSSessionInfo,
            WTSConfigInfo,
            WTSValidationInfo,
            WTSSessionAddressV4,
            WTSIsRemoteSession
        }

        public class WTSSession
        {
            public int SessionId { get; set; }
            public WTS_CONNECTSTATE_CLASS State { get; set; }
            public string? Workstation { get; set; }
            public string? IPAddress { get; set; }
            public string? Username { get; set; }
            public int HorizontalResolution { get; set; }
            public int VerticalResolution { get; set; }
            public int ColorDepth { get; set; }
            public string? ClientApplicationDirectory { get; set; }
        }

        public static IEnumerable<WTSSession> GetSessions()
        {
            var sessions = new List<WTSSession>();

            var pServer = nint.Zero;
            var pSessionInfo = nint.Zero;

            try
            {
                var count = 0;
                var sessionCount = WTSEnumerateSessions(pServer, 0, 1, ref pSessionInfo, ref count);
                var dataSize = (long)Marshal.SizeOf(typeof(WTS_SESSION_INFO));
                var current = (long)pSessionInfo;

                if (sessionCount <= 0)
                    return sessions;

                for (int i = 0; i < count; i++)
                {
                    if (Marshal.PtrToStructure((nint)current, typeof(WTS_SESSION_INFO)) is not object sessionStructure)
                        continue;

                    var sessionInfo = (WTS_SESSION_INFO)sessionStructure;
                    current += dataSize;

                    var session = new WTSSession
                    {
                        SessionId = sessionInfo.iSessionID,
                        State = sessionInfo.oState,
                        Workstation = sessionInfo.sWinsWorkstationName,
                    };

                    var returned = (uint)0;

                    // get terminal user address
                    var address = nint.Zero;
                    var clientAddress = new WTS_CLIENT_ADDRESS();

                    if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSClientAddress, out address, out returned) == true)
                    {
                        if (Marshal.PtrToStructure(address, clientAddress.GetType()) is not object addressStructure)
                            break;

                        clientAddress = (WTS_CLIENT_ADDRESS)addressStructure;
                        session.IPAddress = clientAddress.bAddress[2] + "." + clientAddress.bAddress[3] + "." + clientAddress.bAddress[4] + "." + clientAddress.bAddress[5];
                    }

                    // get terminal user name
                    if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSUserName, out address, out returned) == true)
                    {
                        session.Username = Marshal.PtrToStringAnsi(address);
                    }

                    // get terminal user domain name
                    if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSDomainName, out address, out returned) == true)
                    {
                        session.Username = Marshal.PtrToStringAnsi(address) + @"\" + session.Username;
                    }

                    // get terminal user display informations
                    var clientDisplay = new WTS_CLIENT_DISPLAY();

                    if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSClientDisplay, out address, out returned) == true)
                    {
                        if (Marshal.PtrToStructure(address, clientDisplay.GetType()) is not object displayStructure)
                            break;

                        clientDisplay = (WTS_CLIENT_DISPLAY)displayStructure;
                        session.HorizontalResolution = clientDisplay.iHorizontalResolution;
                        session.VerticalResolution = clientDisplay.iVerticalResolution;
                        session.ColorDepth = clientDisplay.iColorDepth;
                    }

                    // get terminal user application directory
                    if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSClientDirectory, out address, out returned) == true)
                    {
                        session.ClientApplicationDirectory = Marshal.PtrToStringAnsi(address);
                    }

                    sessions.Add(session);
                }
            }
            catch (Exception)
            {
                throw;
            }
            finally
            {
                WTSFreeMemory(pSessionInfo);
            }

            return sessions;
        }
    }
}