Merge pull request #233 from dev-kittens/feature/delete-bot-responses

Allow deleting bot responses with ❌ reaction

PluralKit.Bot/Bot.cs

@@ -9,6 +9,8 @@ using App.Metrics;
 
 using Autofac;
 
+using Dapper;
+
 using DSharpPlus;
 using DSharpPlus.Entities;
 using DSharpPlus.EventArgs;
@@ -34,18 +36,21 @@ namespace PluralKit.Bot
         private readonly PeriodicStatCollector _collector;
         private readonly IMetrics _metrics;
         private readonly ErrorMessageService _errorMessageService;
+        private readonly IDatabase _db;
 
         private bool _hasReceivedReady = false;
         private Timer _periodicTask; // Never read, just kept here for GC reasons
 
-        public Bot(DiscordShardedClient client, ILifetimeScope services, ILogger logger, PeriodicStatCollector collector, IMetrics metrics, ErrorMessageService errorMessageService)
+        public Bot(DiscordShardedClient client, ILifetimeScope services, ILogger logger, PeriodicStatCollector collector, IMetrics metrics, 
+            ErrorMessageService errorMessageService, IDatabase db)
         {
             _client = client;
+            _logger = logger.ForContext<Bot>();
             _services = services;
             _collector = collector;
             _metrics = metrics;
             _errorMessageService = errorMessageService;
-            _logger = logger.ForContext<Bot>();
+            _db = db;
         }
 
         public void Init()
@@ -177,6 +182,9 @@ namespace PluralKit.Bot
 
             await UpdateBotStatus();
 
+            // Clean up message cache in postgres
+            await _db.Execute(conn => conn.QueryAsync("select from cleanup_command_message()"));
+
             // Collect some stats, submit them to the metrics backend
             await _collector.CollectStats();
             await Task.WhenAll(((IMetricsRoot) _metrics).ReportRunner.RunAllAsync());

PluralKit.Bot/CommandSystem/Context.cs

@@ -64,7 +64,7 @@ namespace PluralKit.Bot
         internal IDatabase Database => _db;
         internal ModelRepository Repository => _repo;
 
-        public Task<DiscordMessage> Reply(string text = null, DiscordEmbed embed = null, IEnumerable<IMention> mentions = null)
+        public async Task<DiscordMessage> Reply(string text = null, DiscordEmbed embed = null, IEnumerable<IMention> mentions = null)
         {
             if (!this.BotHasAllPermissions(Permissions.SendMessages))
                 // Will be "swallowed" during the error handler anyway, this message is never shown.
@@ -72,7 +72,12 @@ namespace PluralKit.Bot
 
             if (embed != null && !this.BotHasAllPermissions(Permissions.EmbedLinks))
                 throw new PKError("PluralKit does not have permission to send embeds in this channel. Please ensure I have the **Embed Links** permission enabled.");
-            return Channel.SendMessageFixedAsync(text, embed: embed, mentions: mentions);
+            var msg = await Channel.SendMessageFixedAsync(text, embed: embed, mentions: mentions);
+            if (embed != null) 
+                // Sensitive information that might want to be deleted by :x: reaction is typically in an embed format (member cards, for example)
+                // This may need to be changed at some point but works well enough for now
+                await _db.Execute(conn => _repo.SaveCommandMessage(conn, msg.Id, Author.Id));
+            return msg;
         }
         
         public async Task Execute<T>(Command commandDef, Func<T, Task> handler)

PluralKit.Bot/Commands/CommandTree.cs

@@ -189,9 +189,9 @@ namespace PluralKit.Bot
             if (ctx.Match("random", "r"))
                 return ctx.Execute<Member>(MemberRandom, m => m.MemberRandom(ctx));
 
-            ctx.Reply(
+            // remove compiler warning
+            return ctx.Reply(
                 $"{Emojis.Error} Unknown command {ctx.PeekArgument().AsCode()}. For a list of possible commands, see <https://pluralkit.me/commands>.");
-            return Task.CompletedTask;
         }
 
         private async Task HandleSystemCommand(Context ctx)

PluralKit.Bot/Handlers/ReactionAdded.cs

@@ -48,12 +48,15 @@ namespace PluralKit.Bot
                 _db.Execute(c => _repo.GetMessage(c, evt.Message.Id));
 
             FullMessage msg;
+            CommandMessage cmdmsg;
             switch (evt.Emoji.Name)
             {
                 // Message deletion
                 case "\u274C": // Red X
                     if ((msg = await GetMessage()) != null)
                         await HandleDeleteReaction(evt, msg);
+                    else if ((cmdmsg = await _db.Execute(conn => _repo.GetCommandMessage(conn, evt.Message.Id))) != null)
+                        await HandleCommandDeleteReaction(evt, cmdmsg);
                     break;                
                 
                 case "\u2753": // Red question mark
@@ -92,6 +95,25 @@ namespace PluralKit.Bot
             await _db.Execute(c => _repo.DeleteMessage(c, evt.Message.Id));
         }
 
+        private async ValueTask HandleCommandDeleteReaction(MessageReactionAddEventArgs evt, CommandMessage msg)
+        {
+            if (!evt.Channel.BotHasAllPermissions(Permissions.ManageMessages)) return;
+
+            // Can only delete your own message
+            if (msg.author_id != evt.User.Id) return;
+
+            try
+            {
+                await evt.Message.DeleteAsync();
+            }
+            catch (NotFoundException)
+            {
+                // Message was deleted by something/someone else before we got to it
+            }
+
+            // No need to delete database row here, it'll get deleted by the once-per-minute scheduled task.
+        }
+
         private async ValueTask HandleQueryReaction(MessageReactionAddEventArgs evt, FullMessage msg)
         {
             // Try to DM the user info about the message

PluralKit.Core/Database/Database.cs

@@ -19,7 +19,7 @@ namespace PluralKit.Core
     internal class Database: IDatabase
     {
         private const string RootPath = "PluralKit.Core.Database"; // "resource path" root for SQL files
-        private const int TargetSchemaVersion = 10;
+        private const int TargetSchemaVersion = 11;
         
         private readonly CoreConfig _config;
         private readonly ILogger _logger;

PluralKit.Core/Database/Migrations/11.sql

@@ -0,0 +1,17 @@
+-- SCHEMA VERSION 11: (insert date) --
+-- Create command message table --
+
+create table command_message
+(
+	message_id bigint primary key,
+	author_id bigint not null,
+	timestamp timestamp not null default now()
+);
+
+create function cleanup_command_message() returns void as $$
+begin
+    delete from command_message where timestamp < now() - interval '2 hours';
+end;
+$$ language plpgsql;
+
+update info set schema_version = 11;

PluralKit.Core/Database/Repository/ModelRepository.CommandMessage.cs

@@ -0,0 +1,24 @@
+using System.Collections.Generic;
+using System.Data;
+using System.Threading.Tasks;
+
+using Dapper;
+
+namespace PluralKit.Core
+{
+    public partial class ModelRepository
+	{
+		public Task SaveCommandMessage(IPKConnection conn, ulong message_id, ulong author_id) =>
+			conn.QueryAsync("insert into command_message (message_id, author_id) values (@Message, @Author)",
+				new {Message = message_id, Author = author_id });
+
+		public Task<CommandMessage> GetCommandMessage(IPKConnection conn, ulong message_id) =>
+			conn.QuerySingleOrDefaultAsync<CommandMessage>("select message_id, author_id from command_message where message_id = @Message", 
+				new {Message = message_id});
+	}
+
+	public class CommandMessage
+	{
+		public ulong author_id { get; set; }
+	}
+}