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

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

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

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

        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;
        }
    }
}