77 lines
3.2 KiB
C#
77 lines
3.2 KiB
C#
|
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<AgentSession>
|
|||
|
|
{
|
|||
|
|
private readonly IMongoDatabase _database;
|
|||
|
|
|
|||
|
|
public SessionHandler(IMongoDatabase database)
|
|||
|
|
{
|
|||
|
|
_database = database;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
public async ValueTask HandleAsync<TMessage>(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<Session> sessions, 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<HostSessionEntity>>();
|
|||
|
|
|
|||
|
|
if (sessions is not null && sessions.Any())
|
|||
|
|
{
|
|||
|
|
foreach (var sess in sessions)
|
|||
|
|
{
|
|||
|
|
var filterDefinition = Builders<HostSessionEntity>.Filter.And(new List<FilterDefinition<HostSessionEntity>>
|
|||
|
|
{
|
|||
|
|
Builders<HostSessionEntity>.Filter.Eq(x => x.Host, hostEntity.Id),
|
|||
|
|
Builders<HostSessionEntity>.Filter.Eq(x => x.Sid, sess.Sid)
|
|||
|
|
});
|
|||
|
|
|
|||
|
|
var updateDefinition = Builders<HostSessionEntity>.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<HostSessionEntity>(filterDefinition, updateDefinition)
|
|||
|
|
{
|
|||
|
|
IsUpsert = true
|
|||
|
|
});
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
bulk.Add(new DeleteManyModel<HostSessionEntity>(Builders<HostSessionEntity>.Filter.And(new List<FilterDefinition<HostSessionEntity>>
|
|||
|
|
{
|
|||
|
|
Builders<HostSessionEntity>.Filter.Eq(x => x.Host, hostEntity.Id),
|
|||
|
|
Builders<HostSessionEntity>.Filter.Ne(x => x.Batch, batch)
|
|||
|
|
})));
|
|||
|
|
|
|||
|
|
var result = await _database.HostSession().BulkWriteAsync(bulk, cancellationToken: cancellationToken);
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|