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

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

namespace Insight.Agent.Network.Handlers;

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

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

    private static List<Drive> GetDrives()
    {
        using var searcher = new ManagementObjectSearcher
        {
            Scope = new ManagementScope(@"root\cimv2"),
            Query = new ObjectQuery("select index, name, caption, model, manufacturer, serialNumber, size, status, interfacetype, firmwarerevision, deviceid, pnpdeviceid from win32_diskdrive")
        };

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

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

        var drives = new List<Drive>();

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

                var properties = @object.GetPropertyHashes();

                drive.Index = @object.GetValue<uint>(properties, "index");
                drive.Id = @object.GetValue<string>(properties, "deviceid")?.Trim();
                drive.Name = @object.GetValue<string>(properties, "model")?.Trim();
                drive.Manufacturer = @object.GetValue<string>(properties, "manufacturer")?.Trim();
                drive.SerialNumber = @object.GetValue<string>(properties, "serialnumber")?.Trim();
                drive.Size = @object.GetValue<ulong>(properties, "size");
                drive.Status = @object.GetValue<string>(properties, "status")?.Trim();
                drive.InterfaceType = @object.GetValue<string>(properties, "interfacetype")?.Trim();
                drive.FirmwareRevision = @object.GetValue<string>(properties, "firmwarerevision")?.Trim();
                drive.PNPDeviceID = @object.GetValue<string>(properties, "pnpdeviceid")?.Trim();
                drive.Volumes = new List<Volume>();

                var diskpartition = @object.GetRelated("win32_diskpartition");
                using (diskpartition)
                {
                    foreach (ManagementObject dp in diskpartition.Cast<ManagementObject>())
                    {
                        var volume = new Volume();
                        var dpProperties = dp.GetPropertyHashes();

                        volume.NumberOfBlocks = dp.GetValue<ulong>(dpProperties, "numberofblocks");
                        volume.BootPartition = dp.GetValue<bool>(dpProperties, "bootpartition");
                        volume.PrimaryPartition = dp.GetValue<bool>(dpProperties, "primarypartition");
                        volume.Size = dp.GetValue<ulong>(dpProperties, "size");
                        volume.Index = dp.GetValue<uint>(dpProperties, "index");
                        volume.Type = dp.GetValue<string>(dpProperties, "type")?.Trim();
                        volume.Bootable = dp.GetValue<bool>(dpProperties, "bootable");
                        volume.BlockSize = dp.GetValue<ulong>(dpProperties, "blocksize");
                        volume.StartingOffset = dp.GetValue<ulong>(dpProperties, "startingoffset");

                        var logicaldisk = dp.GetRelated("win32_logicaldisk");
                        using (logicaldisk)
                        {
                            foreach (ManagementObject ld in logicaldisk.Cast<ManagementObject>())
                            {
                                var ldProperties = ld.GetPropertyHashes();

                                volume.Id = ld.GetValue<string>(ldProperties, "deviceid")?.Trim();
                                volume.Name = ld.GetValue<string>(ldProperties, "volumename")?.Trim();
                                volume.SerialNumber = ld.GetValue<string>(ldProperties, "volumeserialnumber")?.Trim();
                                volume.DriveType = (DriveType)ld.GetValue<uint>(ldProperties, "drivetype");
                                volume.FileSystem = ld.GetValue<string>(ldProperties, "filesystem")?.Trim();
                                volume.Compressed = ld.GetValue<bool>(ldProperties, "compressed");
                                volume.Size = ld.GetValue<ulong>(ldProperties, "size");
                                volume.FreeSpace = ld.GetValue<ulong>(ldProperties, "freespace");
                                volume.ProviderName = ld.GetValue<string>(ldProperties, "providername")?.Trim();
                            }
                        }

                        drive.Volumes.Add(volume);
                    }
                }

                drives.Add(drive);
            }
        }

        return drives;
    }

    private static List<Volume> GetVolumes()
    {
        using var searcher = new ManagementObjectSearcher
        {
            Scope = new ManagementScope(@"root\cimv2"),
            Query = new ObjectQuery("select deviceid, volumename, volumeserialnumber, drivetype, filesystem, compressed, size, freeSpace, providername from win32_logicaldisk")
        };

        // per device query
        // "ASSOCIATORS OF {Win32_DiskDrive.DeviceID='" + driveDeviceId + "'} WHERE AssocClass=Win32_DiskDriveToDiskPartition"

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

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

        var volumes = new List<Volume>();

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

                var properties = @object.GetPropertyHashes();

                //volume.DeviceId = @object.GetValue<string>(properties, "deviceid")?.Trim();
                //volume.VolumeName = @object.GetValue<string>(properties, "volumename")?.Trim();
                //volume.VolumeSerialNumber = @object.GetValue<string>(properties, "volumeserialnumber")?.Trim();
                volume.DriveType = (DriveType)@object.GetValue<uint>(properties, "drivetype");
                volume.FileSystem = @object.GetValue<string>(properties, "filesystem")?.Trim();
                volume.Compressed = @object.GetValue<bool>(properties, "compressed");
                volume.Size = @object.GetValue<ulong>(properties, "size");
                volume.FreeSpace = @object.GetValue<ulong>(properties, "freespace");
                volume.ProviderName = @object.GetValue<string>(properties, "providername")?.Trim();

                if (volume.Id is not null)
                {
                    searcher.Query = new ObjectQuery("associators of {win32_logicaldisk.deviceid='" + volume.Id + "'} where assocclass=win32_logicaldisktopartition");

                    if (searcher.TryGet(out var collection2))
                    {
                        using (collection2)
                        {
                            foreach (ManagementObject @object2 in collection2)
                            {
                                var properties2 = @object2.GetPropertyHashes();

                                volume.Index = @object2.GetValue<uint>(properties2, "index");
                                //volume.DiskIndex = @object2.GetValue<uint>(properties2, "diskindex");
                                volume.Type = @object2.GetValue<string>(properties2, "type")?.Trim();
                                volume.Bootable = @object2.GetValue<bool>(properties2, "bootable");
                                volume.PrimaryPartition = @object2.GetValue<bool>(properties2, "primarypartition");
                                volume.BootPartition = @object2.GetValue<bool>(properties2, "bootpartition");
                                volume.BlockSize = @object2.GetValue<ulong>(properties2, "blocksize");
                                volume.NumberOfBlocks = @object2.GetValue<ulong>(properties2, "numberofblocks");
                                volume.StartingOffset = @object2.GetValue<ulong>(properties2, "startingoffset");
                            }
                        }
                    }
                }

                volumes.Add(volume);
            }
        }

        return volumes;
    }
}