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

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<AgentSession>
    {
        private readonly IMongoDatabase _database;

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

        public async ValueTask HandleAsync<TMessage>(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<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 bulk = new List<WriteModel<HostUpdateEntity>>();

            if (updates is not null)
            {
                if (updates.Installed is not null && updates.Installed.Any())
                {
                    foreach (var update in updates.Installed)
                    {
                        var filterDefinition = Builders<HostUpdateEntity>.Filter.And(new List<FilterDefinition<HostUpdateEntity>>
                            {
                                Builders<HostUpdateEntity>.Filter.Eq(x => x.Host, hostEntity.Id),
                                Builders<HostUpdateEntity>.Filter.Eq(x => x.Serial, update.Id)
                            });

                        var updateDefinition = Builders<HostUpdateEntity>.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<HostUpdateEntity>(filterDefinition, updateDefinition)
                        {
                            IsUpsert = true
                        });
                    }
                }

                if (updates.Pending is not null && updates.Pending.Any())
                {
                    foreach (var update in updates.Pending)
                    {
                        var filterDefinition = Builders<HostUpdateEntity>.Filter.And(new List<FilterDefinition<HostUpdateEntity>>
                            {
                                Builders<HostUpdateEntity>.Filter.Eq(x => x.Host, hostEntity.Id),
                                Builders<HostUpdateEntity>.Filter.Eq(x => x.Serial, update.Id)
                            });

                        var updateDefinition = Builders<HostUpdateEntity>.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<HostUpdateEntity>(filterDefinition, updateDefinition)
                        {
                            IsUpsert = true
                        });
                    }
                }
            }

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

            var result = await _database.HostUpdate().BulkWriteAsync(bulk, cancellationToken: cancellationToken);
        }
    }
}