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

using Insight.Domain.Interfaces;
using Insight.Domain.Network;
using Insight.Domain.Network.Agent.Messages;
using Insight.Infrastructure.Entities;
using MongoDB.Bson;
using MongoDB.Driver;

namespace Insight.Server.Network.Agent.Handlers;

public class DriveHandler : IMessageHandler<AgentSession>
{
    private readonly IMongoDatabase _database;

    public DriveHandler(IMongoDatabase database)
    {
        _database = database;
    }

    public async ValueTask HandleAsync<TMessage>(AgentSession sender, TMessage message, CancellationToken cancellationToken) where TMessage : IMessage
    {
        switch (message)
        {
            case Collection<Drive> drives:
                await OnDrivesAsync(sender, drives, cancellationToken);
                break;
        }
    }

    private async ValueTask OnDrivesAsync(AgentSession session, List<Drive> drives, CancellationToken cancellationToken)
    {
        var agentEntity = await _database.Agent().Find(Builders<AgentEntity>.Filter.Eq(p => p.Id, session.Id)).FirstOrDefaultAsync(cancellationToken);
        if (agentEntity is null) return;

        var hostEntity = await _database.Host().Find(Builders<HostEntity>.Filter.Eq(p => p.Agent, agentEntity.Id)).FirstOrDefaultAsync(cancellationToken);
        if (hostEntity is null) return;

        var batch = ObjectId.GenerateNewId().ToString();
        var date = DateTime.Now;

        var driveBulk = new List<WriteModel<HostDriveEntity>>();

        if (drives is not null && drives.Any())
        {
            foreach (var drive in drives)
            {
                var driveFilter = Builders<HostDriveEntity>.Filter.And(new List<FilterDefinition<HostDriveEntity>>
                    {
                        Builders<HostDriveEntity>.Filter.Eq(x => x.Host, hostEntity.Id),
                        Builders<HostDriveEntity>.Filter.Eq(x => x.Index, drive.Index)
                    });

                var driveUpdate = Builders<HostDriveEntity>.Update
                    .SetOnInsert(p => p.Insert, date)
                    .SetOnInsert(p => p.Host, hostEntity.Id)
                    .SetOnInsert(p => p.Index, drive.Index)
                    .Set(p => p.Update, date)
                    .Set(p => p.Batch, batch)
                    .Set(p => p.Company, drive.Manufacturer)
                    .Set(p => p.Name, drive.Name)
                    .Set(p => p.Size, drive.Size)
                    .Set(p => p.Type, drive.InterfaceType)
                    .Set(p => p.Serial, drive.SerialNumber)
                    .Set(p => p.Firmware, drive.FirmwareRevision)
                    .Set(p => p.Status, drive.Status)
                    .Set(p => p.Pnp, drive.PNPDeviceID);

                driveBulk.Add(new UpdateOneModel<HostDriveEntity>(driveFilter, driveUpdate)
                {
                    IsUpsert = true
                });
            }
        }

        driveBulk.Add(new DeleteManyModel<HostDriveEntity>(Builders<HostDriveEntity>.Filter.And(new List<FilterDefinition<HostDriveEntity>>
                {
                    Builders<HostDriveEntity>.Filter.Eq(x => x.Host, hostEntity.Id),
                    Builders<HostDriveEntity>.Filter.Ne(x => x.Batch, batch)
                })));

        var driveResult = await _database.HostDrive().BulkWriteAsync(driveBulk, cancellationToken: cancellationToken);

        // volumes

        var volumeBulk = new List<WriteModel<HostVolumeEntity>>();

        if (drives is not null && drives.Any())
        {
            foreach (var drive in drives)
            {
                var driveId = await _database.HostDrive()
                    .Find(p => p.Host == hostEntity.Id && p.Index == drive.Index)
                    .Project(p => p.Id)
                    .FirstOrDefaultAsync();

                if (drive.Volumes is not null && drive.Volumes.Any())
                {
                    foreach (var volume in drive.Volumes)
                    {
                        var volumeFilter = Builders<HostVolumeEntity>.Filter.And(new List<FilterDefinition<HostVolumeEntity>>
                            {
                                Builders<HostVolumeEntity>.Filter.Eq(x => x.Host, hostEntity.Id),
                                Builders<HostVolumeEntity>.Filter.Eq(x => x.Drive, driveId),
                                Builders<HostVolumeEntity>.Filter.Eq(x => x.Index, volume.Index)
                            });

                        var volumeUpdate = Builders<HostVolumeEntity>.Update
                            .SetOnInsert(p => p.Insert, date)
                            .SetOnInsert(p => p.Host, hostEntity.Id)
                            .SetOnInsert(p => p.Drive, driveId)
                            .SetOnInsert(p => p.Index, volume.Index)
                            .Set(p => p.Update, date)
                            .Set(p => p.Batch, batch)
                            .Set(p => p.Name, volume.Name)
                            .Set(p => p.Label, volume.Id)
                            .Set(p => p.Serial, volume.SerialNumber)
                            .Set(p => p.Size, volume.Size)
                            .Set(p => p.FreeSpace, volume.FreeSpace)
                            .Set(p => p.Type, volume.Type)
                            .Set(p => p.FileSystem, volume.FileSystem)
                            .Set(p => p.Compressed, volume.Compressed)
                            .Set(p => p.Bootable, volume.Bootable)
                            .Set(p => p.Primary, volume.PrimaryPartition)
                            .Set(p => p.Boot, volume.Bootable)
                            .Set(p => p.BlockSize, volume.BlockSize)
                            .Set(p => p.Blocks, volume.NumberOfBlocks)
                            .Set(p => p.StartingOffset, volume.StartingOffset)
                            .Set(p => p.Provider, volume.ProviderName);

                        volumeBulk.Add(new UpdateOneModel<HostVolumeEntity>(volumeFilter, volumeUpdate)
                        {
                            IsUpsert = true
                        });
                    }
                }
            }
        }

        volumeBulk.Add(new DeleteManyModel<HostVolumeEntity>(Builders<HostVolumeEntity>.Filter.And(new List<FilterDefinition<HostVolumeEntity>>
                        {
                            Builders<HostVolumeEntity>.Filter.Eq(x => x.Host, hostEntity.Id),
                            Builders<HostVolumeEntity>.Filter.Ne(x => x.Batch, batch)
                        })));

        var volumeResult = await _database.HostVolume().BulkWriteAsync(volumeBulk, cancellationToken: cancellationToken);
    }
}