Refactor system/member avatar code

PluralKit.Bot/CommandSystem/ContextArgumentsExt.cs

@@ -60,6 +60,9 @@ namespace PluralKit.Bot
             return potentialMatches.Any(potentialMatch => flags.Contains(potentialMatch));
         }
 
+        public static bool MatchClear(this Context ctx) =>
+            ctx.Match("clear", "remove", "reset") || ctx.MatchFlag("c", "clear");
+
         public static async Task<List<PKMember>> ParseMemberList(this Context ctx, SystemId? restrictToSystem)
         {
             var members = new List<PKMember>();

PluralKit.Bot/CommandSystem/ContextEntityArgumentsExt.cs

@@ -2,7 +2,6 @@
 
 using DSharpPlus;
 using DSharpPlus.Entities;
-using DSharpPlus.Exceptions;
 
 using PluralKit.Bot.Utils;
 using PluralKit.Core;

PluralKit.Bot/Commands/Avatars/ContextAvatarExt.cs

@@ -0,0 +1,71 @@
+#nullable enable
+using System;
+using System.Linq;
+using System.Text.RegularExpressions;
+using System.Threading.Tasks;
+
+using DSharpPlus;
+using DSharpPlus.Entities;
+
+namespace PluralKit.Bot
+{
+    public static class ContextAvatarExt
+    {
+        // Rewrite cdn.discordapp.com URLs to media.discordapp.net for jpg/png files
+        // This lets us add resizing parameters to "borrow" their media proxy server to downsize the image
+        // which in turn makes it more likely to be underneath the size limit!
+        private static readonly Regex DiscordCdnUrl = new Regex(@"^https?://(?:cdn\.discordapp\.com|media\.discordapp\.net)/attachments/(\d{17,19})/(\d{17,19})/([^\.?/]+)\.(png|jpg|jpeg).*");
+        private static readonly string DiscordMediaUrlReplacement = "https://media.discordapp.net/attachments/$1/$2/$3.$4?width=256&height=256";
+        
+        public static async Task<ParsedImage?> MatchImage(this Context ctx)
+        {
+            // If we have a user @mention/ID, use their avatar 
+            if (await ctx.MatchUser() is { } user)
+            {
+                var url = user.GetAvatarUrl(ImageFormat.Png, 256);
+                return new ParsedImage {Url = url, Source = AvatarSource.User, SourceUser = user};
+            }
+            
+            // If we have a positional argument, try to parse it as a URL
+            var arg = ctx.RemainderOrNull();
+            if (arg != null)
+            {
+                if (!Uri.TryCreate(arg, UriKind.Absolute, out var uri))
+                    throw Errors.InvalidUrl(arg);
+
+                if (uri.Scheme != "http" && uri.Scheme != "https")
+                    throw Errors.InvalidUrl(arg);
+
+                return new ParsedImage {Url = TryRewriteCdnUrl(uri.ToString()), Source = AvatarSource.Url};
+            }
+            
+            // If we have an attachment, use that 
+            if (ctx.Message.Attachments.FirstOrDefault() is {} attachment)
+            {
+                var url = TryRewriteCdnUrl(attachment.ProxyUrl);
+                return new ParsedImage {Url = url, Source = AvatarSource.Attachment};
+            }
+            
+            // We should only get here if there are no arguments (which would get parsed as URL + throw if error)
+            // and if there are no attachments (which would have been caught just before)
+            return null;
+        }
+
+        private static string TryRewriteCdnUrl(string url) =>
+            DiscordCdnUrl.Replace(url, DiscordMediaUrlReplacement);
+    }
+
+    public struct ParsedImage
+    {
+        public string Url;
+        public AvatarSource Source;
+        public DiscordUser? SourceUser;
+    }
+
+    public enum AvatarSource
+    {
+        Url,
+        User,
+        Attachment
+    }
+}

PluralKit.Bot/Commands/CommandTree.cs

@@ -15,7 +15,7 @@ namespace PluralKit.Bot
         public static Command SystemRename = new Command("system name", "system rename [name]", "Renames your system");
         public static Command SystemDesc = new Command("system description", "system description [description]", "Changes your system's description");
         public static Command SystemTag = new Command("system tag", "system tag [tag]", "Changes your system's tag");
-        public static Command SystemAvatar = new Command("system avatar", "system avatar [url|@mention]", "Changes your system's avatar");
+        public static Command SystemAvatar = new Command("system icon", "system icon [url|@mention]", "Changes your system's icon");
         public static Command SystemDelete = new Command("system delete", "system delete", "Deletes your system");
         public static Command SystemTimezone = new Command("system timezone", "system timezone [timezone]", "Changes your system's time zone");
         public static Command SystemProxy = new Command("system proxy", "system proxy [on|off]", "Enables or disables message proxying in a specific server");

PluralKit.Bot/Commands/MemberAvatar.cs

@@ -1,11 +1,7 @@
 #nullable enable
 using System;
-using System.Linq;
 using System.Threading.Tasks;
 
-using Dapper;
-
-using DSharpPlus;
 using DSharpPlus.Entities;
 
 using PluralKit.Core;
@@ -23,7 +19,6 @@ namespace PluralKit.Bot
         
         private async Task AvatarClear(AvatarLocation location, Context ctx, PKMember target, MemberGuildSettings? mgs)
         {
-            ctx.CheckSystem().CheckOwnMember(target);
             await UpdateAvatar(location, ctx, target, null);
             if (location == AvatarLocation.Server)
             {
@@ -43,9 +38,6 @@ namespace PluralKit.Bot
 
         private async Task AvatarShow(AvatarLocation location, Context ctx, PKMember target, MemberGuildSettings? guildData)
         {
-            var field = location == AvatarLocation.Server ? $"server avatar (for {ctx.Guild.Name})" : "avatar";
-            var cmd = location == AvatarLocation.Server ? "serveravatar" : "avatar";
-            
             var currentValue = location == AvatarLocation.Member ? target.AvatarUrl : guildData?.AvatarUrl;
             var canAccess = location != AvatarLocation.Member || target.AvatarPrivacy.CanAccess(ctx.LookupContextFor(target));
             if (string.IsNullOrEmpty(currentValue) || !canAccess)
@@ -61,6 +53,9 @@ namespace PluralKit.Bot
                     throw new PKError($"This member does not have a server avatar set. Type `pk;member {target.Hid} avatar` to see their global avatar.");
             }
             
+            var field = location == AvatarLocation.Server ? $"server avatar (for {ctx.Guild.Name})" : "avatar";
+            var cmd = location == AvatarLocation.Server ? "serveravatar" : "avatar";
+            
             var eb = new DiscordEmbedBuilder()
                 .WithTitle($"{target.NameFor(ctx)}'s {field}")
                 .WithImageUrl(currentValue);
@@ -69,47 +64,6 @@ namespace PluralKit.Bot
             await ctx.Reply(embed: eb.Build());
         }
 
-        private async Task AvatarFromUser(AvatarLocation location, Context ctx, PKMember target, DiscordUser user)
-        {
-            ctx.CheckSystem().CheckOwnMember(target);
-            if (user.AvatarHash == null) throw Errors.UserHasNoAvatar;
-            
-            var url = user.GetAvatarUrl(ImageFormat.Png, 256);
-            await UpdateAvatar(location, ctx, target, url);
-            
-            var embed = new DiscordEmbedBuilder().WithImageUrl(url).Build();
-            if (location == AvatarLocation.Server)
-                await ctx.Reply($"{Emojis.Success} Member server avatar changed to {user.Username}'s avatar! This avatar will now be used when proxying in this server (**{ctx.Guild.Name}**). {Emojis.Warn} Please note that if {user.Username} changes their avatar, the member's server avatar will need to be re-set.", embed: embed);
-            else if (location == AvatarLocation.Member)
-                await ctx.Reply($"{Emojis.Success} Member avatar changed to {user.Username}'s avatar! {Emojis.Warn} Please note that if {user.Username} changes their avatar, the member's avatar will need to be re-set.", embed: embed);
-        }
-
-        private async Task AvatarFromArg(AvatarLocation location, Context ctx, PKMember target, string url)
-        {
-            ctx.CheckSystem().CheckOwnMember(target);
-            if (url.Length > Limits.MaxUriLength) throw Errors.InvalidUrl(url);
-            await AvatarUtils.VerifyAvatarOrThrow(url);
-
-            await UpdateAvatar(location, ctx, target, url);
-
-            var embed = new DiscordEmbedBuilder().WithImageUrl(url).Build();
-            if (location == AvatarLocation.Server)
-                await ctx.Reply($"{Emojis.Success} Member server avatar changed. This avatar will now be used when proxying in this server (**{ctx.Guild.Name}**).", embed: embed);
-            else
-                await ctx.Reply($"{Emojis.Success} Member avatar changed.");
-        }
-
-        private async Task AvatarFromAttachment(AvatarLocation location, Context ctx, PKMember target, DiscordAttachment attachment)
-        {
-            ctx.CheckSystem().CheckOwnMember(target);
-            await AvatarUtils.VerifyAvatarOrThrow(attachment.Url);
-            await UpdateAvatar(location, ctx, target, attachment.Url);
-            if (location == AvatarLocation.Server)
-                await ctx.Reply($"{Emojis.Success} Member server avatar changed to attached image. This avatar will now be used when proxying in this server (**{ctx.Guild.Name}**). Please note that if you delete the message containing the attachment, the avatar will stop working.");
-            else if (location == AvatarLocation.Member)
-                await ctx.Reply($"{Emojis.Success} Member avatar changed to attached image. Please note that if you delete the message containing the attachment, the avatar will stop working.");
-        }
-
         public async Task ServerAvatar(Context ctx, PKMember target)
         {
             ctx.CheckGuildContext();
@@ -128,28 +82,77 @@ namespace PluralKit.Bot
 
         private async Task AvatarCommandTree(AvatarLocation location, Context ctx, PKMember target, MemberGuildSettings? guildData)
         {
-            if (ctx.Match("clear", "remove", "reset") || ctx.MatchFlag("c", "clear"))
+            // First, see if we need to *clear*
+            if (ctx.MatchClear())
+            {
+                ctx.CheckSystem().CheckOwnMember(target);
                 await AvatarClear(location, ctx, target, guildData);
-            else if (ctx.RemainderOrNull() == null && ctx.Message.Attachments.Count == 0)
+                return;
-                await AvatarShow(location, ctx, target, guildData);
-            else if (await ctx.MatchUser() is {} user)
-                await AvatarFromUser(location, ctx, target, user);
-            else if (ctx.RemainderOrNull() is {} url)
-                await AvatarFromArg(location, ctx, target, url);
-            else if (ctx.Message.Attachments.FirstOrDefault() is {} attachment)
-                await AvatarFromAttachment(location, ctx, target, attachment);
-            else throw new Exception("Unexpected condition when parsing avatar command");
             }
 
-        private Task UpdateAvatar(AvatarLocation location, Context ctx, PKMember target, string? avatar)
+            // Then, parse an image from the command (from various sources...)
+            var avatarArg = await ctx.MatchImage();
+            if (avatarArg == null)
+            {
+                // If we didn't get any, just show the current avatar
+                await AvatarShow(location, ctx, target, guildData);
+                return;
+            }
+
+            ctx.CheckSystem().CheckOwnMember(target);
+            await ValidateUrl(avatarArg.Value.Url);
+            await UpdateAvatar(location, ctx, target, avatarArg.Value.Url);
+            await PrintResponse(location, ctx, target, avatarArg.Value, guildData);
+        }
+
+        private static Task ValidateUrl(string url)
+        {
+            if (url.Length > Limits.MaxUriLength)
+                throw Errors.InvalidUrl(url);
+            return AvatarUtils.VerifyAvatarOrThrow(url);
+        }
+
+        private Task PrintResponse(AvatarLocation location, Context ctx, PKMember target, ParsedImage avatar,
+                                   MemberGuildSettings? targetGuildData)
+        {
+            var typeFrag = location switch
+            {
+                AvatarLocation.Server => "server avatar",
+                AvatarLocation.Member => "avatar",
+                _ => throw new ArgumentOutOfRangeException(nameof(location))
+            };
+            
+            var serverFrag = location switch
+            {
+                AvatarLocation.Server => $" This avatar will now be used when proxying in this server (**{ctx.Guild.Name}**).",
+                AvatarLocation.Member when targetGuildData?.AvatarUrl != null => $"\n{Emojis.Note} Note that this member *also* has a server-specific avatar set in this server (**{ctx.Guild.Name}**), and thus changing the global avatar will have no effect here.",
+                _ => ""
+            };
+
+            var msg = avatar.Source switch
+            {
+                AvatarSource.User => $"{Emojis.Success} Member {typeFrag} changed to {avatar.SourceUser?.Username}'s avatar!{serverFrag}\n{Emojis.Warn} If {avatar.SourceUser?.Username} changes their avatar, the member's avatar will need to be re-set.",
+                AvatarSource.Url => $"{Emojis.Success} Member {typeFrag} changed to the image at the given URL.{serverFrag}",
+                AvatarSource.Attachment => $"{Emojis.Success} Member {typeFrag} changed to attached image.{serverFrag}\n{Emojis.Warn} If you delete the message containing the attachment, the avatar will stop working.",
+                _ => throw new ArgumentOutOfRangeException()
+            };
+
+            // The attachment's already right there, no need to preview it.
+            var hasEmbed = avatar.Source != AvatarSource.Attachment;
+            return hasEmbed 
+                ? ctx.Reply(msg, embed: new DiscordEmbedBuilder().WithImageUrl(avatar.Url).Build()) 
+                : ctx.Reply(msg);
+        }
+
+        private Task UpdateAvatar(AvatarLocation location, Context ctx, PKMember target, string? url)
         {
             switch (location)
             {
                 case AvatarLocation.Server:
-                    var serverPatch = new MemberGuildPatch { AvatarUrl = avatar };
+                    var serverPatch = new MemberGuildPatch { AvatarUrl = url };
                     return _db.Execute(c => c.UpsertMemberGuild(target.Id, ctx.Guild.Id, serverPatch));
                 case AvatarLocation.Member:
-                    var memberPatch = new MemberPatch { AvatarUrl = avatar };
+                    var memberPatch = new MemberPatch { AvatarUrl = url };
                     return _db.Execute(c => c.UpdateMember(target.Id, memberPatch));
                 default:
                     throw new ArgumentOutOfRangeException($"Unknown avatar location {location}");

PluralKit.Bot/Commands/SystemEdit.cs

@@ -132,56 +132,55 @@ namespace PluralKit.Bot
         {
             ctx.CheckSystem();
 
-            if (ctx.Match("clear") || ctx.MatchFlag("c", "clear"))
+            async Task ClearIcon()
             {
-                var patch = new SystemPatch {AvatarUrl = null};
+                await _db.Execute(c => c.UpdateSystem(ctx.System.Id, new SystemPatch {AvatarUrl = null}));
-                await _db.Execute(conn => conn.UpdateSystem(ctx.System.Id, patch));
+                await ctx.Reply($"{Emojis.Success} System icon cleared.");
-                
-                await ctx.Reply($"{Emojis.Success} System avatar cleared.");
-                return;
             }
-            else if (ctx.RemainderOrNull() == null && ctx.Message.Attachments.Count == 0)
+
+            async Task SetIcon(ParsedImage img)
+            {
+                if (img.Url.Length > Limits.MaxUriLength) 
+                    throw Errors.InvalidUrl(img.Url);
+                await AvatarUtils.VerifyAvatarOrThrow(img.Url);
+
+                await _db.Execute(c => c.UpdateSystem(ctx.System.Id, new SystemPatch {AvatarUrl = img.Url}));
+            
+                var msg = img.Source switch
+                {
+                    AvatarSource.User => $"{Emojis.Success} System icon changed to {img.SourceUser?.Username}'s avatar!\n{Emojis.Warn} If {img.SourceUser?.Username} changes their avatar, the system icon will need to be re-set.",
+                    AvatarSource.Url => $"{Emojis.Success} System icon changed to the image at the given URL.",
+                    AvatarSource.Attachment => $"{Emojis.Success} System icon changed to attached image.\n{Emojis.Warn} If you delete the message containing the attachment, the system icon will stop working.",
+                    _ => throw new ArgumentOutOfRangeException()
+                };
+            
+                // The attachment's already right there, no need to preview it.
+                var hasEmbed = img.Source != AvatarSource.Attachment;
+                await (hasEmbed 
+                    ? ctx.Reply(msg, embed: new DiscordEmbedBuilder().WithImageUrl(img.Url).Build()) 
+                    : ctx.Reply(msg));
+            }
+
+            async Task ShowIcon()
             {
                 if ((ctx.System.AvatarUrl?.Trim() ?? "").Length > 0)
                 {
                     var eb = new DiscordEmbedBuilder()
-                        .WithTitle($"System avatar")
+                        .WithTitle("System icon")
                         .WithImageUrl(ctx.System.AvatarUrl)
-                        .WithDescription($"To clear, use `pk;system avatar clear`.");
+                        .WithDescription("To clear, use `pk;system icon clear`.");
                     await ctx.Reply(embed: eb.Build());
                 }
                 else
-                    throw new PKSyntaxError($"This system does not have an avatar set. Set one by attaching an image to this command, or by passing an image URL or @mention.");
+                    throw new PKSyntaxError("This system does not have an icon set. Set one by attaching an image to this command, or by passing an image URL or @mention.");
-
-                return;
             }
 
-            var member = await ctx.MatchUser();
+            if (ctx.MatchClear())
-            if (member != null)
+                await ClearIcon();
-            {
+            else if (await ctx.MatchImage() is {} img)
-                if (member.AvatarHash == null) throw Errors.UserHasNoAvatar;
+                await SetIcon(img);
-
-                var newUrl = member.GetAvatarUrl(ImageFormat.Png, size: 256);
-                var patch = new SystemPatch {AvatarUrl = newUrl};
-                await _db.Execute(conn => conn.UpdateSystem(ctx.System.Id, patch));
-                
-                var embed = new DiscordEmbedBuilder().WithImageUrl(newUrl).Build();
-                await ctx.Reply(
-                    $"{Emojis.Success} System avatar changed to {member.Username}'s avatar! {Emojis.Warn} Please note that if {member.Username} changes their avatar, the system's avatar will need to be re-set.", embed: embed);
-            }
             else
-            {
+                await ShowIcon();
-                // They can't both be null - otherwise we would've hit the conditional at the very top
-                string url = ctx.RemainderOrNull() ?? ctx.Message.Attachments.FirstOrDefault()?.ProxyUrl;
-                if (url?.Length > Limits.MaxUriLength) throw Errors.InvalidUrl(url);
-                await ctx.BusyIndicator(() => AvatarUtils.VerifyAvatarOrThrow(url));
-
-                var patch = new SystemPatch {AvatarUrl = url};
-                await _db.Execute(conn => conn.UpdateSystem(ctx.System.Id, patch));
-
-                var embed = url != null ? new DiscordEmbedBuilder().WithImageUrl(url).Build() : null;
-                await ctx.Reply($"{Emojis.Success} System avatar changed.", embed: embed);
-            }
         }
         
         public async Task Delete(Context ctx) {