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

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

namespace Insight.Agent.Network.Handlers;

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

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

    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> 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.Cast<ManagementObject>())
            {
                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.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;
    }
}