From 6f7ffda63bca4d07f848d0e7e89e583d4dd65119 Mon Sep 17 00:00:00 2001 From: Noikoio Date: Sun, 17 Feb 2019 15:59:02 -0800 Subject: [PATCH] Ported EntryAutoRole from RegexBot Untested and missing some error handling at the moment. --- Kerobot/Common/EntityName.cs | 32 +++++ Modules-PublicInstance/EntryRole/EntryRole.cs | 127 ++++++++++++++++++ Modules-PublicInstance/EntryRole/GuildData.cs | 77 +++++++++++ 3 files changed, 236 insertions(+) create mode 100644 Modules-PublicInstance/EntryRole/EntryRole.cs create mode 100644 Modules-PublicInstance/EntryRole/GuildData.cs diff --git a/Kerobot/Common/EntityName.cs b/Kerobot/Common/EntityName.cs index b0c4e00..24edd92 100644 --- a/Kerobot/Common/EntityName.cs +++ b/Kerobot/Common/EntityName.cs @@ -2,6 +2,7 @@ using Discord.WebSocket; using System; using System.Collections.Generic; +using System.Linq; namespace Kerobot.Common { @@ -117,6 +118,7 @@ namespace Kerobot.Common /// If known, outputs the ID of the corresponding entity. /// Specifies if the internal ID value should be stored if a match is found. /// True if the ID is known. + [Obsolete] public bool TryResolve(SocketGuild searchGuild, out ulong id, bool keepId, EntityType searchType) { if (Id.HasValue) @@ -215,5 +217,35 @@ namespace Kerobot.Common else return $"{pf}{Name}"; } + + #region Helper methods + /// + /// Attempts to find the corresponding role within the given guild. + /// + /// The guild in which to search for the role. + /// + /// Specifies if this EntityName instance should keep the snowflake ID of the + /// corresponding role found in this guild, if it is not already known by this instance. + /// + /// + public SocketRole FindRoleIn(SocketGuild guild, bool updateMissingID = false) + { + if (this.Type != EntityType.Role) + throw new ArgumentException("This EntityName instance must correspond to a Role."); + + bool dirty = false; // flag to update ID if possible regardless of updateMissingID setting + if (this.Id.HasValue) + { + var role = guild.GetRole(Id.Value); + if (role != null) return role; + else dirty = true; // only set if ID already existed but is now invalid + } + + var r = guild.Roles.FirstOrDefault(rq => string.Equals(rq.Name, this.Name, StringComparison.OrdinalIgnoreCase)); + if (r != null && (updateMissingID || dirty)) this.Id = r.Id; + + return r; + } + #endregion } } diff --git a/Modules-PublicInstance/EntryRole/EntryRole.cs b/Modules-PublicInstance/EntryRole/EntryRole.cs new file mode 100644 index 0000000..b380e17 --- /dev/null +++ b/Modules-PublicInstance/EntryRole/EntryRole.cs @@ -0,0 +1,127 @@ +using Discord.WebSocket; +using Newtonsoft.Json.Linq; +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; + +namespace Kerobot.Modules.EntryRole +{ + /// + /// Automatically sets a role onto users entering the guild. + /// + // TODO add persistent role support, make it an option + [KerobotModule] + public class EntryRole : ModuleBase + { + readonly Task _workerTask; + readonly CancellationTokenSource _workerTaskToken; // TODO make use of this when possible + + public EntryRole(Kerobot kb) : base(kb) + { + DiscordClient.UserJoined += DiscordClient_UserJoined; + DiscordClient.UserLeft += DiscordClient_UserLeft; + + _workerTaskToken = new CancellationTokenSource(); + _workerTask = Task.Factory.StartNew(RoleApplyWorker, _workerTaskToken.Token, + TaskCreationOptions.LongRunning, TaskScheduler.Default); + } + + private Task DiscordClient_UserJoined(SocketGuildUser arg) + { + GetGuildState(arg.Guild.Id)?.WaitlistAdd(arg.Id); + return Task.CompletedTask; + } + + private Task DiscordClient_UserLeft(SocketGuildUser arg) + { + GetGuildState(arg.Guild.Id)?.WaitlistRemove(arg.Id); + return Task.CompletedTask; + } + + public override Task CreateGuildStateAsync(JToken config) + { + if (config == null) return null; + + // preserve previously running timers? + // research: can GetState be called here or is it undefined? + + if (config.Type != JTokenType.Object) + throw new ModuleLoadException("Configuration is not properly defined."); + + return Task.FromResult(new GuildData((JObject)config)); + } + + /// + /// Main worker task. + /// + private async Task RoleApplyWorker() + { + while (!_workerTaskToken.IsCancellationRequested) + { + await Task.Delay(5000); + + var subworkers = new List(); + foreach (var g in DiscordClient.Guilds) + { + subworkers.Add(RoleApplyGuildSubWorker(g)); + } + Task.WaitAll(subworkers.ToArray()); + } + } + + /// + /// Guild-specific processing by worker task. + /// + internal async Task RoleApplyGuildSubWorker(SocketGuild g) + { + var gconf = GetGuildState(g.Id); + if (gconf == null) return; + + // Get users to be affected + ulong[] userIds; + lock (gconf.WaitingList) + { + if (gconf.WaitingList.Count == 0) return; + + var now = DateTimeOffset.UtcNow; + var queryIds = from item in gconf.WaitingList + where item.Value > now + select item.Key; + userIds = queryIds.ToArray(); + + foreach (var item in userIds) gconf.WaitingList.Remove(item); + } + + var gusers = new List(); + foreach (var item in userIds) + { + var gu = g.GetUser(item); + if (gu == null) continue; // silently drop unknown users (is this fine?) + gusers.Add(gu); + } + if (gusers.Count == 0) return; + + // Attempt to get role. + var targetRole = gconf.TargetRole.FindRoleIn(g, true); + if (targetRole == null) + { + // Notify of this failure. + string failList = ""; + foreach (var item in gusers) failList += $", {item.Username}#{item.Discriminator}"; + + await LogAsync(g.Id, "Unable to find role to apply. (Was the role deleted?) " + + "Failed to set role to the following users: " + failList.Substring(2)); + } + + // Apply roles + foreach (var item in gusers) + { + // TODO exception handling and notification on forbidden + if (item.Roles.Contains(targetRole)) continue; + await item.AddRoleAsync(targetRole); + } + } + } +} diff --git a/Modules-PublicInstance/EntryRole/GuildData.cs b/Modules-PublicInstance/EntryRole/GuildData.cs new file mode 100644 index 0000000..7570bd8 --- /dev/null +++ b/Modules-PublicInstance/EntryRole/GuildData.cs @@ -0,0 +1,77 @@ +using Kerobot.Common; +using Newtonsoft.Json.Linq; +using System; +using System.Collections.Generic; + +namespace Kerobot.Modules.EntryRole +{ + /// + /// Contains configuration data as well as per-guild timers for those awaiting role assignment. + /// + class GuildData + { + /// + /// Lock on self. + /// + public Dictionary WaitingList { get; } + + /// + /// Role to apply. + /// + public EntityName TargetRole { get; } + /// + /// Time to wait until applying the role, in seconds. + /// + public int WaitTime { get; } + + const int WaitTimeMax = 600; // 10 minutes + + public GuildData(JObject conf) + { + var cfgRole = conf["Role"]?.Value(); + if (string.IsNullOrWhiteSpace(cfgRole)) + throw new ModuleLoadException("Role value not specified."); + try + { + TargetRole = new EntityName(cfgRole); + } + catch (ArgumentException) + { + throw new ModuleLoadException("Role config value was not properly specified to be a role."); + } + + try + { + WaitTime = conf["WaitTime"].Value(); + } + catch (NullReferenceException) + { + throw new ModuleLoadException("WaitTime value not specified."); + } + catch (InvalidCastException) + { + throw new ModuleLoadException("WaitTime value must be a number."); + } + + if (WaitTime > WaitTimeMax) + { + // don't silently correct it + throw new ModuleLoadException($"WaitTime value may not exceed {WaitTimeMax} seconds."); + } + if (WaitTime < 0) + { + throw new ModuleLoadException("WaitTime value may not be negative."); + } + } + + public void WaitlistAdd(ulong userId) + { + lock (WaitingList) WaitingList.Add(userId, DateTimeOffset.UtcNow.AddSeconds(WaitTime)); + } + + public void WaitlistRemove(ulong userId) + { + lock (WaitingList) WaitingList.Remove(userId); + } + } +}