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 UpdateHandler : IAgentMessageHandler { private readonly IMongoDatabase _database; public UpdateHandler(IMongoDatabase database) { _database = database; } public async ValueTask HandleAsync(AgentSession sender, TMessage message, CancellationToken cancellationToken) where TMessage : IAgentMessage { if (message is UpdateList updates) { await OnUpdatesAsync(sender, updates, cancellationToken); } } private async ValueTask OnUpdatesAsync(AgentSession session, UpdateList updates, 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 bulk = new List>(); if (updates is not null) { if (updates.Installed is not null && updates.Installed.Any()) { foreach (var update in updates.Installed) { var filterDefinition = Builders.Filter.And(new List> { Builders.Filter.Eq(x => x.Host, hostEntity.Id), Builders.Filter.Eq(x => x.Serial, update.Id) }); var updateDefinition = Builders.Update .SetOnInsert(p => p.Insert, date) .SetOnInsert(p => p.Host, hostEntity.Id) .SetOnInsert(p => p.Serial, update.Id) .Set(p => p.Update, date) .Set(p => p.Batch, batch) .Set(p => p.Pending, false) .Set(p => p.Name, update.Name) .Set(p => p.Description, update.Description) .Set(p => p.SupportUrl, update.SupportUrl) .Set(p => p.Result, update.Result.ToString()) .Set(p => p.Date, update.Date); bulk.Add(new UpdateOneModel(filterDefinition, updateDefinition) { IsUpsert = true }); } } if (updates.Pending is not null && updates.Pending.Any()) { foreach (var update in updates.Pending) { var filterDefinition = Builders.Filter.And(new List> { Builders.Filter.Eq(x => x.Host, hostEntity.Id), Builders.Filter.Eq(x => x.Serial, update.Id) }); var updateDefinition = Builders.Update .SetOnInsert(p => p.Insert, date) .SetOnInsert(p => p.Host, hostEntity.Id) .SetOnInsert(p => p.Serial, update.Id) .Set(p => p.Update, date) .Set(p => p.Batch, batch) .Set(p => p.Pending, true) .Set(p => p.Name, update.Name) .Set(p => p.Description, update.Description) .Set(p => p.SupportUrl, update.SupportUrl) .Set(p => p.Result, update.Result?.ToString()) .Set(p => p.Type, update.Type.ToString()) .Set(p => p.Size, update.Size) .Set(p => p.IsDownloaded, update.IsDownloaded) .Set(p => p.CanRequestUserInput, update.CanRequestUserInput) .Set(p => p.RebootBehavior, update.RebootBehavior?.ToString()) .Set(p => p.Date, update.Date); bulk.Add(new UpdateOneModel(filterDefinition, updateDefinition) { IsUpsert = true }); } } } bulk.Add(new DeleteManyModel(Builders.Filter.And(new List> { Builders.Filter.Eq(x => x.Host, hostEntity.Id), Builders.Filter.Ne(x => x.Batch, batch) }))); var result = await _database.HostUpdate().BulkWriteAsync(bulk, cancellationToken: cancellationToken); } } }