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

using Insight.Agent.Interfaces;
using Insight.Agent.Messages;
using System.Management;
using System.Runtime.Versioning;
using static Insight.Agent.Messages.PhysicalDisk;
using static Insight.Agent.Messages.StoragePool;
using static Insight.Agent.Messages.VirtualDisk;

namespace Insight.Agent.Network.Handlers
{
    [SupportedOSPlatform("windows")]
    public class StoragePoolHandler : IAgentMessageHandler<AgentSession>
    {
        public async ValueTask HandleAsync<TMessage>(AgentSession sender, TMessage message, CancellationToken cancellationToken) where TMessage : IAgentMessage
        {
            if (message is GetInventory)
            {
                var result = new StoragePoolList();
                result.AddRange(GetStoragePool());

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

        private static List<StoragePool> GetStoragePool()
        {
            if (Environment.OSVersion.Version.Major < 6 || Environment.OSVersion.Version.Major == 6 && Environment.OSVersion.Version.Minor < 2)
            {
                throw new PlatformNotSupportedException();
            }

            using var searcher = new ManagementObjectSearcher
            {
                Scope = new ManagementScope(@"root\microsoft\windows\storage"),
                Query = new ObjectQuery("select objectid, uniqueid, name, friendlyname, resiliencysettingnamedefault, isprimordial, isreadonly, isclustered, size, allocatedsize, logicalsectorsize, operationalstatus, healthstatus from msft_storagepool")
            };

            if (searcher.TryGet(out var collection) is false)
            {
                searcher.Query = new ObjectQuery("select * from msft_storagepool");

                if (searcher.TryGet(out collection) is false) throw new InvalidOperationException("WMI Collection NULL");
            }

            var pools = new List<StoragePool>();

            using (collection)
            {
                foreach (ManagementObject @object in collection.Cast<ManagementObject>())
                {
                    var pool = new StoragePool();

                    var properties = @object.GetPropertyHashes();

                    pool.UniqueId = @object.GetValue<string>(properties, "uniqueid")?.Trim();
                    pool.Name = @object.GetValue<string>(properties, "name")?.Trim();
                    pool.FriendlyName = @object.GetValue<string>(properties, "friendlyname")?.Trim();

                    if (@object.TryGetValue<ushort[]>(properties, "operationalstatus", out var operationals) && operationals is not null)
                    {
                        pool.States = operationals.Select(p => (StoragePool.OperationalState)p).ToList();
                    }

                    pool.Health = (StoragePool.HealthState)@object.GetValue<ushort>(properties, "healthstatus");
                    pool.RetireMissingPhysicalDisks = (RetireMissingPhysicalDisksEnum)@object.GetValue<ushort>(properties, "retiremissingphysicaldisks");
                    pool.Resiliency = @object.GetValue<string>(properties, "resiliencysettingnamedefault")?.Trim();
                    pool.IsPrimordial = @object.GetValue<bool>(properties, "isprimordial");
                    pool.IsReadOnly = @object.GetValue<bool>(properties, "isreadonly");
                    pool.IsClustered = @object.GetValue<bool>(properties, "isclustered");
                    pool.Size = @object.GetValue<ulong>(properties, "size");
                    pool.AllocatedSize = @object.GetValue<ulong>(properties, "allocatedsize");
                    pool.SectorSize = @object.GetValue<ulong>(properties, "logicalsectorsize");

                    if (@object.GetValue<string>(properties, "objectid") is string objectId)
                    {
                        pool.PhysicalDisks = QueryPhysicalDisksByStoragePool(objectId);
                        pool.VirtualDisks = QueryVirtualDisksByStoragePool(objectId);
                    }

                    pools.Add(pool);
                }
            }

            return pools;
        }

        private static List<PhysicalDisk> GetPhysicalDisks()
        {
            using var searcher = new ManagementObjectSearcher
            {
                Scope = new ManagementScope(@"root\microsoft\windows\storage"),
                Query = new ObjectQuery("select objectid, uniqueid, name, friendlyname from msft_physicaldisk")
            };

            if (searcher.TryGet(out var collection) is false)
            {
                searcher.Query = new ObjectQuery("select * from msft_physicaldisk");

                if (searcher.TryGet(out collection) is false) throw new InvalidOperationException("WMI Collection NULL");
            }

            var disks = new List<PhysicalDisk>();

            using (collection)
            {
                foreach (ManagementObject @object in collection.Cast<ManagementObject>())
                {
                    var disk = new PhysicalDisk();

                    var properties = @object.GetPropertyHashes();

                    disk.UniqueId = @object.GetValue<string>(properties, "uniqueid")?.Trim();
                    disk.FriendlyName = @object.GetValue<string>(properties, "friendlyname")?.Trim();
                    disk.Manufacturer = @object.GetValue<string>(properties, "manufacturer")?.Trim();
                    disk.Model = @object.GetValue<string>(properties, "model")?.Trim();
                    disk.MediaType = @object.GetValue<ushort>(properties, "mediatype");
                    disk.BusType = @object.GetValue<ushort>(properties, "bustype");

                    if (@object.TryGetValue<ushort[]>(properties, "operationalstatus", out var operationals) && operationals is not null)
                    {
                        disk.States = operationals.Select(p => (PhysicalDisk.OperationalState)p).ToList();
                    }

                    disk.Health = (PhysicalDisk.HealthState)@object.GetValue<ushort>(properties, "healthstatus");

                    if (@object.TryGetValue<ushort[]>(properties, "supportedusages", out var supportedusages) && supportedusages is not null)
                    {
                        disk.SupportedUsages = supportedusages.Select(p => (SupportedUsagesEnum)p).ToList();
                    }

                    disk.Usage = @object.GetValue<ushort>(properties, "usage");
                    disk.PhysicalLocation = @object.GetValue<string>(properties, "physicallocation")?.Trim();
                    disk.SerialNumber = @object.GetValue<string>(properties, "serialnumber")?.Trim();
                    disk.FirmwareVersion = @object.GetValue<string>(properties, "firmwareversion")?.Trim();
                    disk.Size = @object.GetValue<ulong>(properties, "size");
                    disk.AllocatedSize = @object.GetValue<ulong>(properties, "allocatedsize");
                    disk.LogicalSectorSize = @object.GetValue<ulong>(properties, "logicalsectorsize");
                    disk.PhysicalSectorSize = @object.GetValue<ulong>(properties, "physicalsectorsize");
                    disk.VirtualDiskFootprint = @object.GetValue<ulong>(properties, "virtualdiskfootprint");

                    disks.Add(disk);
                }
            }

            return disks;
        }

        private static List<VirtualDisk> GetVirtualDisks()
        {
            using var searcher = new ManagementObjectSearcher
            {
                Scope = new ManagementScope(@"root\microsoft\windows\storage"),
                Query = new ObjectQuery("select objectid, uniqueid, name, friendlyname, access, provisioningtype, physicaldiskredundancy, resiliencysettingname, isdeduplicationenabled, issnapshot, operationalstatus, healthstatus, size, allocatedsize, footprintonpool, readcachesize, writecachesize from msft_virtualdisk")
            };

            if (searcher.TryGet(out var collection) is false)
            {
                searcher.Query = new ObjectQuery("select * from msft_virtualdisk");

                if (searcher.TryGet(out collection) is false) throw new InvalidOperationException("WMI Collection NULL");
            }

            var disks = new List<VirtualDisk>();

            using (collection)
            {
                foreach (ManagementObject @object in collection.Cast<ManagementObject>())
                {
                    var disk = new VirtualDisk();

                    var properties = @object.GetPropertyHashes();

                    disk.UniqueId = @object.GetValue<string>(properties, "uniqueid")?.Trim();
                    disk.Name = @object.GetValue<string>(properties, "name")?.Trim();
                    disk.FriendlyName = @object.GetValue<string>(properties, "friendlyname")?.Trim();
                    disk.AccessType = (AccessTypeEnum)@object.GetValue<ushort>(properties, "access");
                    disk.ProvisioningType = (ProvisioningTypeEnum)@object.GetValue<ushort>(properties, "provisioningtype");
                    disk.PhysicalDiskRedundancy = @object.GetValue<ushort>(properties, "physicaldiskredundancy");
                    disk.ResiliencySettingName = @object.GetValue<string>(properties, "resiliencysettingname")?.Trim();
                    disk.Deduplication = @object.GetValue<bool>(properties, "isdeduplicationenabled");
                    disk.IsSnapshot = @object.GetValue<bool>(properties, "issnapshot");

                    if (@object.TryGetValue<ushort[]>(properties, "operationalstatus", out var operationals) && operationals is not null)
                    {
                        disk.States = operationals.Select(p => (VirtualDisk.OperationalState)p).ToList();
                    }

                    disk.Health = (VirtualDisk.HealthState)@object.GetValue<ushort>(properties, "healthstatus");
                    disk.Size = @object.GetValue<ulong>(properties, "size");
                    disk.AllocatedSize = @object.GetValue<ulong>(properties, "allocatedsize");
                    disk.FootprintOnPool = @object.GetValue<ulong>(properties, "footprintonpool");
                    disk.ReadCacheSize = @object.GetValue<ulong>(properties, "readcachesize");
                    disk.WriteCacheSize = @object.GetValue<ulong>(properties, "writecachesize");

                    disks.Add(disk);
                }
            }

            return disks;
        }

        private static List<PhysicalDisk> QueryPhysicalDisksByStoragePool(string storagePoolObjectId)
        {
            using var searcher = new ManagementObjectSearcher
            {
                Scope = new ManagementScope(@"root\microsoft\windows\storage"),
                Query = new ObjectQuery("ASSOCIATORS OF {MSFT_StoragePool.ObjectId=\"" + Helpers.EscapeWql(storagePoolObjectId) + "\"} WHERE AssocClass = MSFT_StoragePoolToPhysicalDisk")
            };

            if (searcher.TryGet(out var collection) is false) throw new InvalidOperationException("WMI Collection NULL");

            var disks = new List<PhysicalDisk>();

            using (collection)
            {
                foreach (ManagementObject @object in collection)
                {
                    var disk = new PhysicalDisk();

                    var properties = @object.GetPropertyHashes();

                    disk.UniqueId = @object.GetValue<string>(properties, "uniqueid")?.Trim();
                    disk.DeviceId = @object.GetValue<string>(properties, "deviceid")?.Trim();
                    disk.FriendlyName = @object.GetValue<string>(properties, "friendlyname")?.Trim();
                    disk.Manufacturer = @object.GetValue<string>(properties, "manufacturer")?.Trim();
                    disk.Model = @object.GetValue<string>(properties, "model")?.Trim();
                    disk.MediaType = @object.GetValue<ushort>(properties, "mediatype");
                    disk.BusType = @object.GetValue<ushort>(properties, "bustype");

                    if (@object.TryGetValue<ushort[]>(properties, "operationalstatus", out var operationals) && operationals is not null)
                    {
                        disk.States = operationals.Select(p => (PhysicalDisk.OperationalState)p).ToList();
                    }

                    disk.Health = (PhysicalDisk.HealthState)@object.GetValue<ushort>(properties, "healthstatus");

                    if (@object.TryGetValue<ushort[]>(properties, "supportedusages", out var supportedusages) && supportedusages is not null)
                    {
                        disk.SupportedUsages = supportedusages.Select(p => (SupportedUsagesEnum)p).ToList();
                    }

                    disk.Usage = @object.GetValue<ushort>(properties, "usage");
                    disk.PhysicalLocation = @object.GetValue<string>(properties, "physicallocation")?.Trim();
                    disk.SerialNumber = @object.GetValue<string>(properties, "serialnumber")?.Trim();
                    disk.FirmwareVersion = @object.GetValue<string>(properties, "firmwareversion")?.Trim();
                    disk.Size = @object.GetValue<ulong>(properties, "size");
                    disk.AllocatedSize = @object.GetValue<ulong>(properties, "allocatedsize");
                    disk.LogicalSectorSize = @object.GetValue<ulong>(properties, "logicalsectorsize");
                    disk.PhysicalSectorSize = @object.GetValue<ulong>(properties, "physicalsectorsize");
                    disk.VirtualDiskFootprint = @object.GetValue<ulong>(properties, "virtualdiskfootprint");

                    disks.Add(disk);
                }
            }

            return disks;
        }

        private static List<VirtualDisk> QueryVirtualDisksByStoragePool(string storagePoolObjectId)
        {
            using var searcher = new ManagementObjectSearcher
            {
                Scope = new ManagementScope(@"root\microsoft\windows\storage"),
                Query = new ObjectQuery("ASSOCIATORS OF {MSFT_StoragePool.ObjectId=\"" + Helpers.EscapeWql(storagePoolObjectId) + "\"} WHERE AssocClass = MSFT_StoragePoolToVirtualDisk")
            };

            if (searcher.TryGet(out var collection) is false) throw new InvalidOperationException("WMI Collection NULL");

            var disks = new List<VirtualDisk>();

            using (collection)
            {
                foreach (ManagementObject @object in collection)
                {
                    var disk = new VirtualDisk();

                    var properties = @object.GetPropertyHashes();

                    disk.UniqueId = @object.GetValue<string>(properties, "uniqueid")?.Trim();
                    disk.Name = @object.GetValue<string>(properties, "name")?.Trim();
                    disk.FriendlyName = @object.GetValue<string>(properties, "friendlyname")?.Trim();
                    disk.AccessType = (AccessTypeEnum)@object.GetValue<ushort>(properties, "access");
                    disk.ProvisioningType = (ProvisioningTypeEnum)@object.GetValue<ushort>(properties, "provisioningtype");
                    disk.PhysicalDiskRedundancy = @object.GetValue<ushort>(properties, "physicaldiskredundancy");
                    disk.ResiliencySettingName = @object.GetValue<string>(properties, "resiliencysettingname")?.Trim();
                    disk.Deduplication = @object.GetValue<bool>(properties, "isdeduplicationenabled");
                    disk.IsSnapshot = @object.GetValue<bool>(properties, "issnapshot");

                    if (@object.TryGetValue<ushort[]>(properties, "operationalstatus", out var operationals) && operationals is not null)
                    {
                        disk.States = operationals.Select(p => (VirtualDisk.OperationalState)p).ToList();
                    }

                    disk.Health = (VirtualDisk.HealthState)@object.GetValue<ushort>(properties, "healthstatus");
                    disk.Size = @object.GetValue<ulong>(properties, "size");
                    disk.AllocatedSize = @object.GetValue<ulong>(properties, "allocatedsize");
                    disk.FootprintOnPool = @object.GetValue<ulong>(properties, "footprintonpool");
                    disk.ReadCacheSize = @object.GetValue<ulong>(properties, "readcachesize");
                    disk.WriteCacheSize = @object.GetValue<ulong>(properties, "writecachesize");

                    disks.Add(disk);
                }
            }

            return disks;
        }
    }
}