using Insight.Agent.Interfaces; using Insight.Agent.Messages; using Insight.Infrastructure; using Insight.Infrastructure.Entities; using MongoDB.Bson; using MongoDB.Driver; namespace Insight.Server.Network.Handlers.Agent { public class DriveHandler : IAgentMessageHandler { private readonly IMongoDatabase _database; public DriveHandler(IMongoDatabase database) { _database = database; } public async ValueTask HandleAsync(AgentSession sender, TMessage message, CancellationToken cancellationToken) where TMessage : IAgentMessage { if (message is DriveList drives) { await OnDrivesAsync(sender, drives, cancellationToken); } } private async ValueTask OnDrivesAsync(AgentSession session, List drives, CancellationToken cancellationToken) { var agentEntity = await _database.Agent().Find(Builders.Filter.Eq(p => p.Id, session.Id)).FirstOrDefaultAsync(cancellationToken); if (agentEntity is null) return; var hostEntity = await _database.Host().Find(Builders.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>(); if (drives is not null && drives.Any()) { foreach (var drive in drives) { var driveFilter = Builders.Filter.And(new List> { Builders.Filter.Eq(x => x.Host, hostEntity.Id), Builders.Filter.Eq(x => x.Index, drive.Index) }); var driveUpdate = Builders.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(driveFilter, driveUpdate) { IsUpsert = true }); } } driveBulk.Add(new DeleteManyModel(Builders.Filter.And(new List> { Builders.Filter.Eq(x => x.Host, hostEntity.Id), Builders.Filter.Ne(x => x.Batch, batch) }))); var driveResult = await _database.HostDrive().BulkWriteAsync(driveBulk, cancellationToken: cancellationToken); // volumes var volumeBulk = new List>(); 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.Filter.And(new List> { Builders.Filter.Eq(x => x.Host, hostEntity.Id), Builders.Filter.Eq(x => x.Drive, driveId), Builders.Filter.Eq(x => x.Index, volume.Index) }); var volumeUpdate = Builders.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(volumeFilter, volumeUpdate) { IsUpsert = true }); } } } } volumeBulk.Add(new DeleteManyModel(Builders.Filter.And(new List> { Builders.Filter.Eq(x => x.Host, hostEntity.Id), Builders.Filter.Ne(x => x.Batch, batch) }))); var volumeResult = await _database.HostVolume().BulkWriteAsync(volumeBulk, cancellationToken: cancellationToken); } } }