#nullable disable using LinqToDB; using LinqToDB.EntityFrameworkCore; using NadekoBot.Common.ModuleBehaviors; using NadekoBot.Db; using NadekoBot.Modules.Patronage; using NadekoBot.Db.Models; using OneOf.Types; using OneOf; namespace NadekoBot.Modules.Administration.Services; public sealed class ReactionRolesService : IReadyExecutor, INService, IReactionRoleService { private readonly DbService _db; private readonly DiscordSocketClient _client; private readonly IBotCredentials _creds; private ConcurrentDictionary> _cache; private readonly object _cacheLock = new(); private readonly SemaphoreSlim _assignementLock = new(1, 1); private readonly IPatronageService _ps; private static readonly FeatureLimitKey _reroFLKey = new() { Key = "rero:max_count", PrettyName = "Reaction Role" }; public ReactionRolesService( DiscordSocketClient client, DbService db, IBotCredentials creds, IPatronageService ps) { _db = db; _ps = ps; _client = client; _creds = creds; _cache = new(); } public async Task OnReadyAsync() { await using var uow = _db.GetDbContext(); var reros = await uow.GetTable() .Where( x => Linq2DbExpressions.GuildOnShard(x.GuildId, _creds.TotalShards, _client.ShardId)) .ToListAsyncLinqToDB(); foreach (var group in reros.GroupBy(x => x.MessageId)) { _cache[group.Key] = group.ToList(); } _client.ReactionAdded += ClientOnReactionAdded; _client.ReactionRemoved += ClientOnReactionRemoved; } private async Task<(IGuildUser, IRole)> GetUserAndRoleAsync( ulong userId, ReactionRoleV2 rero) { var guild = _client.GetGuild(rero.GuildId); var role = guild?.GetRole(rero.RoleId); if (role is null) return default; var user = guild.GetUser(userId) as IGuildUser ?? await _client.Rest.GetGuildUserAsync(guild.Id, userId); if (user is null) return default; return (user, role); } private Task ClientOnReactionRemoved( Cacheable cmsg, Cacheable ch, SocketReaction r) { if (!_cache.TryGetValue(cmsg.Id, out var reros)) return Task.CompletedTask; _ = Task.Run(async () => { var emote = await GetFixedEmoteAsync(cmsg, r.Emote); var rero = reros.FirstOrDefault(x => x.Emote == emote.Name || x.Emote == emote.ToString()); if (rero is null) return; var (user, role) = await GetUserAndRoleAsync(r.UserId, rero); if (user.IsBot) return; await _assignementLock.WaitAsync(); try { if (user.RoleIds.Contains(role.Id)) { await user.RemoveRoleAsync(role.Id, new RequestOptions() { AuditLogReason = $"Reaction role" }); } } finally { _assignementLock.Release(); } }); return Task.CompletedTask; } // had to add this because for some reason, reactionremoved event's reaction doesn't have IsAnimated set, // causing the .ToString() to be wrong on animated custom emotes private async Task GetFixedEmoteAsync( Cacheable cmsg, IEmote inputEmote) { // this should only run for emote if (inputEmote is not Emote e) return inputEmote; // try to get the message and pull var msg = await cmsg.GetOrDownloadAsync(); var emote = msg.Reactions.Keys.FirstOrDefault(x => e.Equals(x)); return emote ?? inputEmote; } private Task ClientOnReactionAdded( Cacheable msg, Cacheable ch, SocketReaction r) { if (!_cache.TryGetValue(msg.Id, out var reros)) return Task.CompletedTask; _ = Task.Run(async () => { var rero = reros.FirstOrDefault(x => x.Emote == r.Emote.Name || x.Emote == r.Emote.ToString()); if (rero is null) return; var (user, role) = await GetUserAndRoleAsync(r.UserId, rero); if (user.IsBot) return; await _assignementLock.WaitAsync(); try { if (!user.RoleIds.Contains(role.Id)) { // first check if there is a level requirement // and if there is, make sure user satisfies it if (rero.LevelReq > 0) { await using var ctx = _db.GetDbContext(); var levelData = await ctx.GetTable() .GetLevelDataFor(user.GuildId, user.Id); if (levelData.Level < rero.LevelReq) return; } // remove all other roles from the same group from the user // execept in group 0, which is a special, non-exclusive group if (rero.Group != 0) { var exclusive = reros .Where(x => x.Group == rero.Group && x.RoleId != role.Id) .Select(x => x.RoleId) .Distinct() .ToArray(); if (exclusive.Any()) { try { await user.RemoveRolesAsync(exclusive, new RequestOptions() { AuditLogReason = "Reaction role exclusive group" }); } catch { } } // remove user's previous reaction try { var m = await msg.GetOrDownloadAsync(); if (m is not null) { var reactToRemove = m.Reactions .FirstOrDefault(x => x.Key.ToString() != r.Emote.ToString()) .Key; if (reactToRemove is not null) { await m.RemoveReactionAsync(reactToRemove, user); } } } catch { } } await user.AddRoleAsync(role.Id, new() { AuditLogReason = "Reaction role" }); } } finally { _assignementLock.Release(); } }); return Task.CompletedTask; } /// /// Adds a single reaction role /// /// Guild where to add a reaction role /// Message to which to add a reaction role /// /// /// /// /// The result of the operation public async Task> AddReactionRole( IGuild guild, IMessage msg, string emote, IRole role, int group = 0, int levelReq = 0) { ArgumentOutOfRangeException.ThrowIfNegative(group); ArgumentOutOfRangeException.ThrowIfNegative(levelReq); await using var ctx = _db.GetDbContext(); await using var tran = await ctx.Database.BeginTransactionAsync(); var activeReactionRoles = await ctx.GetTable() .Where(x => x.GuildId == guild.Id) .CountAsync(); var result = await _ps.TryGetFeatureLimitAsync(_reroFLKey, guild.OwnerId, 50); if (result.Quota != -1 && activeReactionRoles >= result.Quota) return result; await ctx.GetTable() .InsertOrUpdateAsync(() => new() { GuildId = guild.Id, ChannelId = msg.Channel.Id, MessageId = msg.Id, Emote = emote, RoleId = role.Id, Group = group, LevelReq = levelReq }, (old) => new() { RoleId = role.Id, Group = group, LevelReq = levelReq }, () => new() { MessageId = msg.Id, Emote = emote, }); await tran.CommitAsync(); var obj = new ReactionRoleV2() { GuildId = guild.Id, MessageId = msg.Id, Emote = emote, RoleId = role.Id, Group = group, LevelReq = levelReq }; lock (_cacheLock) { _cache.AddOrUpdate(msg.Id, _ => [obj], (_, list) => { list.RemoveAll(x => x.Emote == emote); list.Add(obj); return list; }); } return new Success(); } /// /// Get all reaction roles on the specified server /// /// /// public async Task> GetReactionRolesAsync(ulong guildId) { await using var ctx = _db.GetDbContext(); return await ctx.GetTable() .Where(x => x.GuildId == guildId) .ToListAsync(); } /// /// Remove reaction roles on the specified message /// /// /// /// public async Task RemoveReactionRoles(ulong guildId, ulong messageId) { // guildid is used for quick index lookup await using var ctx = _db.GetDbContext(); var changed = await ctx.GetTable() .Where(x => x.GuildId == guildId && x.MessageId == messageId) .DeleteAsync(); _cache.TryRemove(messageId, out _); if (changed == 0) return false; return true; } /// /// Remove all reaction roles in the specified server /// /// /// public async Task RemoveAllReactionRoles(ulong guildId) { await using var ctx = _db.GetDbContext(); var output = await ctx.GetTable() .Where(x => x.GuildId == guildId) .DeleteWithOutputAsync(x => x.MessageId); lock (_cacheLock) { foreach (var o in output) { _cache.TryRemove(o, out _); } } return output.Length; } public async Task> TransferReactionRolesAsync( ulong guildId, ulong fromMessageId, ulong toMessageId) { await using var ctx = _db.GetDbContext(); var updated = ctx.GetTable() .Where(x => x.GuildId == guildId && x.MessageId == fromMessageId) .UpdateWithOutput(old => new() { MessageId = toMessageId }, (old, neu) => neu); lock (_cacheLock) { if (_cache.TryRemove(fromMessageId, out var data)) { if (_cache.TryGetValue(toMessageId, out var newData)) { newData.AddRange(data); } else { _cache[toMessageId] = data; } } } return updated.Select(x => x.Emote.ToIEmote()).ToList(); } }