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 SessionHandler : IAgentMessageHandler { private readonly IMongoDatabase _database; public SessionHandler(IMongoDatabase database) { _database = database; } public async ValueTask HandleAsync(AgentSession sender, TMessage message, CancellationToken cancellationToken) where TMessage : IAgentMessage { if (message is SessionList sessions) { await OnSessionsAsync(sender, sessions, cancellationToken); } } private async ValueTask OnSessionsAsync(AgentSession session, List sessions, 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 (sessions is not null && sessions.Any()) { foreach (var sess in sessions) { var filterDefinition = Builders.Filter.And(new List> { Builders.Filter.Eq(x => x.Host, hostEntity.Id), Builders.Filter.Eq(x => x.Sid, sess.Sid) }); var updateDefinition = Builders.Update .SetOnInsert(p => p.Insert, date) .SetOnInsert(p => p.Host, hostEntity.Id) .SetOnInsert(p => p.Sid, sess.Sid) .Set(p => p.Update, date) .Set(p => p.Batch, batch) .Set(p => p.User, sess.User) .Set(p => p.Remote, sess.Remote) .Set(p => p.Type, sess.Type) .Set(p => p.State, sess.Status); 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.HostSession().BulkWriteAsync(bulk, cancellationToken: cancellationToken); } } }