Add autoproxy functionality
This commit is contained in:
		| @@ -121,6 +121,7 @@ namespace PluralKit.Bot | ||||
|  | ||||
|             .AddTransient<ProxyCacheService>() | ||||
|             .AddSingleton<WebhookCacheService>() | ||||
|             .AddSingleton<AutoproxyCacheService>() | ||||
|             .AddSingleton<ShardInfoService>() | ||||
|             .AddSingleton<CpuStatService>() | ||||
|  | ||||
|   | ||||
| @@ -11,10 +11,12 @@ namespace PluralKit.Bot.Commands | ||||
|     public class AutoproxyCommands | ||||
|     { | ||||
|         private IDataStore _data; | ||||
|         private AutoproxyCacheService _cache; | ||||
|  | ||||
|         public AutoproxyCommands(IDataStore data) | ||||
|         public AutoproxyCommands(IDataStore data, AutoproxyCacheService cache) | ||||
|         { | ||||
|             _data = data; | ||||
|             _cache = cache; | ||||
|         } | ||||
|  | ||||
|         public async Task Autoproxy(Context ctx) | ||||
| @@ -49,6 +51,7 @@ namespace PluralKit.Bot.Commands | ||||
|                 settings.AutoproxyMode = AutoproxyMode.Off; | ||||
|                 settings.AutoproxyMember = null; | ||||
|                 await _data.SetSystemGuildSettings(ctx.System, ctx.Guild.Id, settings); | ||||
|                 await _cache.FlushCacheForSystem(ctx.System, ctx.Guild.Id); | ||||
|                 await ctx.Reply($"{Emojis.Success} Autoproxy turned off in this server."); | ||||
|             } | ||||
|         } | ||||
| @@ -65,6 +68,7 @@ namespace PluralKit.Bot.Commands | ||||
|                 settings.AutoproxyMode = AutoproxyMode.Latch; | ||||
|                 settings.AutoproxyMember = null; | ||||
|                 await _data.SetSystemGuildSettings(ctx.System, ctx.Guild.Id, settings); | ||||
|                 await _cache.FlushCacheForSystem(ctx.System, ctx.Guild.Id); | ||||
|                 await ctx.Reply($"{Emojis.Success} Autoproxy set to latch mode in this server. Messages will now be autoproxied using the *last-proxied member* in this server."); | ||||
|             } | ||||
|         } | ||||
| @@ -81,6 +85,7 @@ namespace PluralKit.Bot.Commands | ||||
|                 settings.AutoproxyMode = AutoproxyMode.Front; | ||||
|                 settings.AutoproxyMember = null; | ||||
|                 await _data.SetSystemGuildSettings(ctx.System, ctx.Guild.Id, settings); | ||||
|                 await _cache.FlushCacheForSystem(ctx.System, ctx.Guild.Id); | ||||
|                 await ctx.Reply($"{Emojis.Success} Autoproxy set to front mode in this server. Messages will now be autoproxied using the *current first fronter*, if any."); | ||||
|             } | ||||
|         } | ||||
| @@ -93,7 +98,7 @@ namespace PluralKit.Bot.Commands | ||||
|             settings.AutoproxyMode = AutoproxyMode.Member; | ||||
|             settings.AutoproxyMember = member.Id; | ||||
|             await _data.SetSystemGuildSettings(ctx.System, ctx.Guild.Id, settings); | ||||
|  | ||||
|             await _cache.FlushCacheForSystem(ctx.System, ctx.Guild.Id); | ||||
|             await ctx.Reply($"{Emojis.Success} Autoproxy set to **{member.Name}** in this server."); | ||||
|         } | ||||
|  | ||||
|   | ||||
							
								
								
									
										73
									
								
								PluralKit.Bot/Services/AutoproxyCacheService.cs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										73
									
								
								PluralKit.Bot/Services/AutoproxyCacheService.cs
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,73 @@ | ||||
| using System; | ||||
| using System.Linq; | ||||
| using System.Threading.Tasks; | ||||
|  | ||||
| using Dapper; | ||||
|  | ||||
| using Microsoft.Extensions.Caching.Memory; | ||||
|  | ||||
| namespace PluralKit.Bot | ||||
| { | ||||
|     public class AutoproxyCacheResult | ||||
|     { | ||||
|         public SystemGuildSettings GuildSettings; | ||||
|         public PKSystem System; | ||||
|         public PKMember AutoproxyMember; | ||||
|     } | ||||
|     public class AutoproxyCacheService | ||||
|     { | ||||
|         private IMemoryCache _cache; | ||||
|         private IDataStore _data; | ||||
|         private DbConnectionFactory _conn; | ||||
|  | ||||
|         public AutoproxyCacheService(IMemoryCache cache, DbConnectionFactory conn, IDataStore data) | ||||
|         { | ||||
|             _cache = cache; | ||||
|             _conn = conn; | ||||
|             _data = data; | ||||
|         } | ||||
|  | ||||
|         public async Task<AutoproxyCacheResult> GetGuildSettings(ulong account, ulong guild) =>  | ||||
|             await _cache.GetOrCreateAsync(GetKey(account, guild), entry => FetchSettings(account, guild, entry)); | ||||
|  | ||||
|         public async Task FlushCacheForSystem(PKSystem system, ulong guild) | ||||
|         { | ||||
|             foreach (var account in await _data.GetSystemAccounts(system)) | ||||
|                 FlushCacheFor(account, guild); | ||||
|         } | ||||
|  | ||||
|         public void FlushCacheFor(ulong account, ulong guild) =>  | ||||
|             _cache.Remove(GetKey(account, guild)); | ||||
|  | ||||
|         private async Task<AutoproxyCacheResult> FetchSettings(ulong account, ulong guild, ICacheEntry entry) | ||||
|         { | ||||
|             using var conn = await _conn.Obtain(); | ||||
|             var data = (await conn.QueryAsync<SystemGuildSettings, PKSystem, PKMember, AutoproxyCacheResult>( | ||||
|                 "select system_guild.*, systems.*, members.* from accounts inner join systems on systems.id = accounts.system inner join system_guild on system_guild.system = systems.id left join members on system_guild.autoproxy_member = members.id where accounts.uid = @Uid and system_guild.guild = @Guild", | ||||
|                 (guildSettings, system, autoproxyMember) => new AutoproxyCacheResult | ||||
|                 { | ||||
|                     GuildSettings = guildSettings, | ||||
|                     System = system, | ||||
|                     AutoproxyMember = autoproxyMember | ||||
|                 }, | ||||
|                 new {Uid = account, Guild = guild})).FirstOrDefault(); | ||||
|              | ||||
|             if (data != null) | ||||
|             { | ||||
|                 // Long expiry for accounts with no system/settings registered | ||||
|                 entry.SetSlidingExpiration(TimeSpan.FromMinutes(5)); | ||||
|                 entry.SetAbsoluteExpiration(TimeSpan.FromHours(1)); | ||||
|             } | ||||
|             else | ||||
|             { | ||||
|                 // Shorter expiry if they already have settings | ||||
|                 entry.SetSlidingExpiration(TimeSpan.FromMinutes(1)); | ||||
|                 entry.SetAbsoluteExpiration(TimeSpan.FromMinutes(5)); | ||||
|             } | ||||
|              | ||||
|             return data; | ||||
|         } | ||||
|  | ||||
|         private string GetKey(ulong account, ulong guild) => $"_system_guild_{account}_{guild}"; | ||||
|     } | ||||
| } | ||||
| @@ -9,6 +9,9 @@ using Discord; | ||||
| using Discord.Net; | ||||
| using Discord.WebSocket; | ||||
|  | ||||
| using NodaTime; | ||||
| using NodaTime.Extensions; | ||||
|  | ||||
| using PluralKit.Core; | ||||
|  | ||||
| using Serilog; | ||||
| @@ -18,7 +21,7 @@ namespace PluralKit.Bot | ||||
|     class ProxyMatch { | ||||
|         public PKMember Member; | ||||
|         public PKSystem System; | ||||
|         public ProxyTag ProxyTags; | ||||
|         public ProxyTag? ProxyTags; | ||||
|         public string InnerText; | ||||
|     } | ||||
|  | ||||
| @@ -31,8 +34,9 @@ namespace PluralKit.Bot | ||||
|         private ILogger _logger; | ||||
|         private WebhookExecutorService _webhookExecutor; | ||||
|         private ProxyCacheService _cache; | ||||
|         private AutoproxyCacheService _autoproxyCache; | ||||
|          | ||||
|         public ProxyService(IDiscordClient client, LogChannelService logChannel, IDataStore data, EmbedService embeds, ILogger logger, ProxyCacheService cache, WebhookExecutorService webhookExecutor, DbConnectionFactory conn) | ||||
|         public ProxyService(IDiscordClient client, LogChannelService logChannel, IDataStore data, EmbedService embeds, ILogger logger, ProxyCacheService cache, WebhookExecutorService webhookExecutor, DbConnectionFactory conn, AutoproxyCacheService autoproxyCache) | ||||
|         { | ||||
|             _client = client; | ||||
|             _logChannel = logChannel; | ||||
| @@ -41,6 +45,7 @@ namespace PluralKit.Bot | ||||
|             _cache = cache; | ||||
|             _webhookExecutor = webhookExecutor; | ||||
|             _conn = conn; | ||||
|             _autoproxyCache = autoproxyCache; | ||||
|             _logger = logger.ForContext<ProxyService>(); | ||||
|         } | ||||
|  | ||||
| @@ -92,8 +97,13 @@ namespace PluralKit.Bot | ||||
|             if (!(message.Channel is ITextChannel channel)) return; | ||||
|              | ||||
|             // Find a member with proxy tags matching the message | ||||
|             var results = await _cache.GetResultsFor(message.Author.Id); | ||||
|             var results = (await _cache.GetResultsFor(message.Author.Id)).ToList(); | ||||
|             var match = GetProxyTagMatch(message.Content, results); | ||||
|              | ||||
|             // If we didn't get a match by proxy tags, try to get one by autoproxy | ||||
|             if (match == null) match = await GetAutoproxyMatch(message, channel); | ||||
|              | ||||
|             // If we still haven't found any, just yeet | ||||
|             if (match == null) return; | ||||
|  | ||||
|             // Gather all "extra" data from DB at once | ||||
| @@ -122,8 +132,9 @@ namespace PluralKit.Bot | ||||
|             if (proxyName.Length > Limits.MaxProxyNameLength) throw Errors.ProxyNameTooLong(proxyName); | ||||
|  | ||||
|             // Add the proxy tags into the proxied message if that option is enabled | ||||
|             var messageContents = match.Member.KeepProxy | ||||
|                 ? $"{match.ProxyTags.Prefix}{match.InnerText}{match.ProxyTags.Suffix}" | ||||
|             // Also check if the member has any proxy tags - some cases autoproxy can return a member with no tags | ||||
|             var messageContents = (match.Member.KeepProxy && match.ProxyTags.HasValue) | ||||
|                 ? $"{match.ProxyTags.Value.Prefix}{match.InnerText}{match.ProxyTags.Value.Suffix}" | ||||
|                 : match.InnerText; | ||||
|              | ||||
|             // Sanitize @everyone, but only if the original user wouldn't have permission to | ||||
| @@ -138,7 +149,7 @@ namespace PluralKit.Bot | ||||
|             ); | ||||
|  | ||||
|             // Store the message in the database, and log it in the log channel (if applicable) | ||||
|             await _data.AddMessage(message.Author.Id, hookMessageId, message.Channel.Id, message.Id, match.Member); | ||||
|             await _data.AddMessage(message.Author.Id, hookMessageId, channel.GuildId, message.Channel.Id, message.Id, match.Member); | ||||
|             await _logChannel.LogMessage(match.System, match.Member, hookMessageId, message.Id, message.Channel as IGuildChannel, message.Author, match.InnerText, aux.Guild); | ||||
|  | ||||
|             // Wait a second or so before deleting the original message | ||||
| @@ -155,6 +166,57 @@ namespace PluralKit.Bot | ||||
|             } | ||||
|         } | ||||
|  | ||||
|         private async Task<ProxyMatch> GetAutoproxyMatch(IMessage message, IGuildChannel channel) | ||||
|         { | ||||
|             // For now we use a backslash as an "escape character", subject to change later | ||||
|             if ((message.Content ?? "").TrimStart().StartsWith("\"")) return null;  | ||||
|              | ||||
|             // Fetch info from the cache, bail if we don't have anything (either no system or no autoproxy settings - AP defaults to off so this works) | ||||
|             var autoproxyCache = await _autoproxyCache.GetGuildSettings(message.Author.Id, channel.GuildId); | ||||
|             if (autoproxyCache == null) return null; | ||||
|  | ||||
|             PKMember member = null; | ||||
|             // Figure out which member to proxy as | ||||
|             switch (autoproxyCache.GuildSettings.AutoproxyMode) | ||||
|             { | ||||
|                 case AutoproxyMode.Off: | ||||
|                     // Autoproxy off, bail | ||||
|                     return null; | ||||
|                 case AutoproxyMode.Front: | ||||
|                     // Front mode: just use the current first fronter | ||||
|                     member = await _data.GetFirstFronter(autoproxyCache.System); | ||||
|                     break; | ||||
|                 case AutoproxyMode.Latch: | ||||
|                     // Latch mode: find last proxied message, use *that* member | ||||
|                     var msg = await _data.GetLastMessageInGuild(message.Author.Id, channel.GuildId); | ||||
|                     if (msg == null) return null; // No message found | ||||
|  | ||||
|                     // If the message is older than 6 hours, ignore it and force the sender to "refresh" a proxy | ||||
|                     // This can be revised in the future, it's a preliminary value. | ||||
|                     var timestamp = SnowflakeUtils.FromSnowflake(msg.Message.Mid).ToInstant(); | ||||
|                     var timeSince = SystemClock.Instance.GetCurrentInstant() - timestamp; | ||||
|                     if (timeSince > Duration.FromHours(6)) return null; | ||||
|                      | ||||
|                     member = msg.Member; | ||||
|                     break; | ||||
|                 case AutoproxyMode.Member: | ||||
|                     // Member mode: just use that member | ||||
|                     member = autoproxyCache.AutoproxyMember; | ||||
|                     break; | ||||
|             } | ||||
|  | ||||
|             // If we haven't found the member (eg. front mode w/ no fronter), bail again | ||||
|             if (member == null) return null; | ||||
|             return new ProxyMatch | ||||
|             { | ||||
|                 System = autoproxyCache.System, | ||||
|                 Member = member, | ||||
|                 // Autoproxying members with no proxy tags is possible, return the correct result | ||||
|                 ProxyTags = member.ProxyTags.Count > 0 ? member.ProxyTags.First() : (ProxyTag?) null, | ||||
|                 InnerText = message.Content | ||||
|             }; | ||||
|         } | ||||
|  | ||||
|         private static string SanitizeEveryoneMaybe(IMessage message, string messageContents) | ||||
|         { | ||||
|             var senderPermissions = ((IGuildUser) message.Author).GetPermissions(message.Channel as IGuildChannel); | ||||
|   | ||||
| @@ -279,12 +279,13 @@ namespace PluralKit { | ||||
|         /// Saves a posted message to the database. | ||||
|         /// </summary> | ||||
|         /// <param name="senderAccount">The ID of the account that sent the original trigger message.</param> | ||||
|         /// <param name="guildId">The ID of the guild the message was posted to.</param> | ||||
|         /// <param name="channelId">The ID of the channel the message was posted to.</param> | ||||
|         /// <param name="postedMessageId">The ID of the message posted by the webhook.</param> | ||||
|         /// <param name="triggerMessageId">The ID of the original trigger message containing the proxy tags.</param> | ||||
|         /// <param name="proxiedMember">The member (and by extension system) that was proxied.</param> | ||||
|         /// <returns></returns> | ||||
|         Task AddMessage(ulong senderAccount, ulong channelId, ulong postedMessageId, ulong triggerMessageId, PKMember proxiedMember); | ||||
|         Task AddMessage(ulong senderAccount, ulong guildId, ulong channelId, ulong postedMessageId, ulong triggerMessageId, PKMember proxiedMember); | ||||
|          | ||||
|         /// <summary> | ||||
|         /// Deletes a message from the data store. | ||||
| @@ -298,6 +299,12 @@ namespace PluralKit { | ||||
|         /// <param name="postedMessageIds">The IDs of the webhook messages to delete.</param> | ||||
|         Task DeleteMessagesBulk(IEnumerable<ulong> postedMessageIds); | ||||
|          | ||||
|         /// <summary> | ||||
|         /// Gets the most recent message sent by a given account in a given guild. | ||||
|         /// </summary> | ||||
|         /// <returns>The full message object, or null if none was found.</returns> | ||||
|         Task<FullMessage> GetLastMessageInGuild(ulong account, ulong guild); | ||||
|          | ||||
|         /// <summary> | ||||
|         /// Gets switches from a system. | ||||
|         /// </summary> | ||||
| @@ -346,6 +353,12 @@ namespace PluralKit { | ||||
|         /// <param name="periodEnd"></param> | ||||
|         /// <returns></returns> | ||||
|         Task<FrontBreakdown> GetFrontBreakdown(PKSystem system, Instant periodStart, Instant periodEnd); | ||||
|  | ||||
|         /// <summary> | ||||
|         /// Gets the first listed fronter in a system. | ||||
|         /// </summary> | ||||
|         /// <returns>The first fronter, or null if none are registered.</returns> | ||||
|         Task<PKMember> GetFirstFronter(PKSystem system); | ||||
|          | ||||
|         /// <summary> | ||||
|         /// Registers a switch with the given members in the given system. | ||||
| @@ -675,10 +688,11 @@ namespace PluralKit { | ||||
|             using (var conn = await _conn.Obtain()) | ||||
|                 return await conn.ExecuteScalarAsync<ulong>("select count(id) from members"); | ||||
|         } | ||||
|         public async Task AddMessage(ulong senderId, ulong messageId, ulong channelId, ulong originalMessage, PKMember member) { | ||||
|         public async Task AddMessage(ulong senderId, ulong messageId, ulong guildId, ulong channelId, ulong originalMessage, PKMember member) { | ||||
|             using (var conn = await _conn.Obtain()) | ||||
|                 await conn.ExecuteAsync("insert into messages(mid, channel, member, sender, original_mid) values(@MessageId, @ChannelId, @MemberId, @SenderId, @OriginalMid)", new { | ||||
|                 await conn.ExecuteAsync("insert into messages(mid, guild, channel, member, sender, original_mid) values(@MessageId, @GuildId, @ChannelId, @MemberId, @SenderId, @OriginalMid)", new { | ||||
|                     MessageId = messageId, | ||||
|                     GuildId = guildId, | ||||
|                     ChannelId = channelId, | ||||
|                     MemberId = member.Id, | ||||
|                     SenderId = senderId, | ||||
| @@ -717,6 +731,17 @@ namespace PluralKit { | ||||
|             } | ||||
|         } | ||||
|  | ||||
|         public async Task<FullMessage> GetLastMessageInGuild(ulong account, ulong guild) | ||||
|         { | ||||
|             using var conn = await _conn.Obtain(); | ||||
|             return (await conn.QueryAsync<PKMessage, PKMember, PKSystem, FullMessage>("select messages.*, members.*, systems.* from messages, members, systems where messages.guild = @Guild and messages.sender = @Uid and messages.member = members.id and systems.id = members.system order by mid desc limit 1", (msg, member, system) => new FullMessage | ||||
|             { | ||||
|                 Message = msg, | ||||
|                 System = system, | ||||
|                 Member = member | ||||
|             }, new { Uid = account, Guild = guild })).FirstOrDefault(); | ||||
|         } | ||||
|  | ||||
|         public async Task<ulong> GetTotalMessages() | ||||
|         { | ||||
|             using (var conn = await _conn.Obtain()) | ||||
| @@ -781,6 +806,15 @@ namespace PluralKit { | ||||
|             }; | ||||
|         } | ||||
|  | ||||
|         public async Task<PKMember> GetFirstFronter(PKSystem system) | ||||
|         { | ||||
|             // TODO: move to extension method since it doesn't rely on internals | ||||
|             var lastSwitch = await GetLatestSwitch(system); | ||||
|             if (lastSwitch == null) return null; | ||||
|  | ||||
|             return await GetSwitchMembers(lastSwitch).FirstOrDefaultAsync(); | ||||
|         } | ||||
|  | ||||
|         public async Task AddSwitch(PKSystem system, IEnumerable<PKMember> members) | ||||
|         { | ||||
|             // Use a transaction here since we're doing multiple executed commands in one | ||||
|   | ||||
		Reference in New Issue
	
	Block a user