From 3b488ab787a053ebdb997d2060f5a7d498d4e732 Mon Sep 17 00:00:00 2001 From: bmintz Date: Sun, 29 Jul 2018 23:34:27 -0500 Subject: [PATCH] add add command a lot of this is from Emoji Connoisseur --- bot.py | 24 +++--- cogs/emoji.py | 198 +++++++++++++++++++++++++++++++++++++++++++++- config.example.py | 2 + utils/__init__.py | 6 ++ utils/emote.py | 19 +++++ utils/errors.py | 45 +++++++++++ utils/misc.py | 44 +++++++++++ 7 files changed, 327 insertions(+), 11 deletions(-) create mode 100644 utils/__init__.py create mode 100644 utils/emote.py create mode 100644 utils/errors.py create mode 100644 utils/misc.py diff --git a/bot.py b/bot.py index 017d319..70815de 100755 --- a/bot.py +++ b/bot.py @@ -4,20 +4,24 @@ import discord from discord.ext import commands -with open('config.py') as f: - config = eval(f.read(), {}) - class Bot(commands.AutoShardedBot): - def __init__(self, **kwargs): - super().__init__(command_prefix=commands.when_mentioned, **kwargs) - for cog in config['cogs']: - self.load_extension(cog) + def __init__(self, **kwargs): + super().__init__(command_prefix=commands.when_mentioned, **kwargs) - async def on_ready(self): - print('Logged on as {0} (ID: {0.id})'.format(self.user)) + with open('config.py') as f: + self.config = eval(f.read(), {}) + + for cog in self.config['cogs']: + self.load_extension(cog) + + def run(self): + super().run(self.config['tokens'].pop('discord')) + + async def on_ready(self): + print('Logged on as {0} (ID: {0.id})'.format(self.user)) bot = Bot() if __name__ == '__main__': - bot.run(config['tokens'].pop('discord')) + bot.run() diff --git a/cogs/emoji.py b/cogs/emoji.py index 0034185..b3d7b4a 100644 --- a/cogs/emoji.py +++ b/cogs/emoji.py @@ -1,16 +1,212 @@ #!/usr/bin/env python3 # encoding: utf-8 +import io +import imghdr +import asyncio +import logging +import traceback +import contextlib + +logger = logging.getLogger(__name__) + +try: + from wand.image import Image +except ImportError: + logger.warn('failed to import wand.image. Image manipulation functions will be unavailable.') + Image = None + +import aiohttp +import discord from discord.ext import commands +import utils +from utils import errors + class Emoji: def __init__(self, bot): self.bot = bot - + self.http = aiohttp.ClientSession(loop=self.bot.loop, read_timeout=30, headers={ + 'User-Agent': + self.bot.config['user_agent'] + ' ' + + self.bot.http.user_agent + }) async def __local_check(self, context): return ( context.guild and context.author.guild_permissions.manage_emojis) + @commands.command() + async def add(self, context, *args): + """Adds an emoji to this server.""" + try: + name, url = self.parse_add_command_args(context, args) + except commands.BadArgument as exception: + return await context.send(exception) + + async with context.typing(): + message = await self.add_safe(context.guild, name, url, context.message.author.id) + await context.send(message) + + @classmethod + def parse_add_command_args(cls, context, args): + if context.message.attachments: + return cls.parse_add_command_attachment(context, args) + + elif len(args) == 1: + match = utils.emote.RE_CUSTOM_EMOTE.match(args[0]) + if match is None: + raise commands.BadArgument( + 'Error: I expected a custom emote as the first argument, ' + 'but I got something else. ' + "If you're trying to add an emote using an image URL, " + 'you need to provide a name as the first argument, like this:\n' + '`{}add NAME_HERE URL_HERE`'.format(context.prefix)) + else: + animated, name, id = match.groups() + url = utils.emote.url(id, animated=animated) + + return name, url + + elif len(args) >= 2: + name = args[0] + match = utils.emote.RE_CUSTOM_EMOTE.match(args[1]) + if match is None: + url = utils.strip_angle_brackets(args[1]) + else: + url = utils.emote.url(match.group('id')) + + return name, url + + elif not args: + raise commands.BadArgument('Your message had no emotes and no name!') + + @staticmethod + def parse_add_command_attachment(context, args): + attachment = context.message.attachments[0] + # as far as i can tell, this is how discord replaces filenames when you upload an emote image + name = ''.join(args) if args else attachment.filename.split('.')[0].replace(' ', '') + url = attachment.url + + return name, url + + async def add_safe(self, guild, name, url, author_id): + """Try to add an emote. Returns a string that should be sent to the user.""" + try: + emote = await self.add_from_url(guild, name, url, author_id) + except discord.HTTPException as ex: + logger.error(traceback.format_exc()) + return ( + 'An error occurred while creating the emote:\n' + + utils.format_http_exception(ex)) + except asyncio.TimeoutError: + return 'Error: retrieving the image took too long.' + except ValueError: + return 'Error: Invalid URL.' + else: + return f'Emote {emote} successfully created.' + + async def add_from_url(self, guild, name, url, author_id): + image_data = await self.fetch_emote(url) + emote = await self.create_emote_from_bytes(guild, name, author_id, image_data) + + return emote + + async def fetch_emote(self, url): + # credits to @Liara#0001 (ID 136900814408122368) for most of this part + # https://gitlab.com/Pandentia/element-zero/blob/47bc8eeeecc7d353ec66e1ef5235adab98ca9635/element_zero/cogs/emoji.py#L217-228 + async with self.http.head(url, timeout=5) as response: + if response.reason != 'OK': + raise errors.HTTPException(response.status) + if response.headers.get('Content-Type') not in ('image/png', 'image/jpeg', 'image/gif'): + raise errors.InvalidImageError + + async with self.http.get(url) as response: + if response.reason != 'OK': + raise errors.HTTPException(response.status) + return io.BytesIO(await response.read()) + + async def create_emote_from_bytes(self, guild, name, author_id, image_data: io.BytesIO): + # resize_until_small is normally blocking, because wand is. + # run_in_executor is magic that makes it non blocking somehow. + # also, None as the executor arg means "use the loop's default executor" + image_data = await self.bot.loop.run_in_executor(None, self.resize_until_small, image_data) + return await guild.create_custom_emoji( + name=name, + image=image_data.read(), + reason=f'Created by {utils.format_user(self.bot, author_id)}') + + @staticmethod + def is_animated(image_data: bytes): + """Return whether the image data is animated, or raise InvalidImageError if it's not an image.""" + type = imghdr.what(None, image_data) + if type == 'gif': + return True + elif type in ('png', 'jpeg'): + return False + else: + raise errors.InvalidImageError + + @classmethod + def size(cls, data: io.BytesIO): + """return the size, in bytes, of the data a BytesIO object represents""" + with cls.preserve_position(data): + data.seek(0, io.SEEK_END) + return data.tell() + + class preserve_position(contextlib.AbstractContextManager): + def __init__(self, fp): + self.fp = fp + self.old_pos = fp.tell() + + def __exit__(self, *excinfo): + self.fp.seek(self.old_pos) + + @classmethod + def resize_until_small(cls, image_data: io.BytesIO): + """If the image_data is bigger than 256KB, resize it until it's not""" + # It's important that we only attempt to resize the image when we have to, + # ie when it exceeds the Discord limit of 256KiB. + # Apparently some <256KiB images become larger when we attempt to resize them, + # so resizing sometimes does more harm than good. + max_resolution = 128 # pixels + size = cls.size(image_data) + while size > 256 * 2**10 and max_resolution >= 32: # don't resize past 32x32 or 256KiB + logger.debug('image size too big (%s bytes)', size) + logger.debug('attempting resize to %s*%s pixels', max_resolution, max_resolution) + image_data = cls.thumbnail(image_data, (max_resolution, max_resolution)) + size = cls.size(image_data) + max_resolution //= 2 + return image_data + + @classmethod + def thumbnail(cls, image_data: io.BytesIO, max_size=(128, 128)): + """Resize an image in place to no more than max_size pixels, preserving aspect ratio.""" + # Credit to @Liara#0001 (ID 136900814408122368) + # https://gitlab.com/Pandentia/element-zero/blob/47bc8eeeecc7d353ec66e1ef5235adab98ca9635/element_zero/cogs/emoji.py#L243-247 + image = Image(blob=image_data) + image.resize(*cls.scale_resolution((image.width, image.height), max_size)) + # we create a new buffer here because there's wand errors otherwise. + # specific error: + # MissingDelegateError: no decode delegate for this image format `' @ error/blob.c/BlobToImage/353 + out = io.BytesIO() + image.save(file=out) + out.seek(0) + return out + + @staticmethod + def scale_resolution(old_res, new_res): + # https://stackoverflow.com/a/6565988 + """Resize a resolution, preserving aspect ratio. Returned w,h will be <= new_res""" + old_width, old_height = old_res + new_width, new_height = new_res + old_ratio = old_width / old_height + new_ratio = new_width / new_height + if new_ratio > old_ratio: + return (old_width * new_height//old_height, new_height) + return new_width, old_height * new_width//old_width + + + def setup(bot): bot.add_cog(Emoji(bot)) diff --git a/config.example.py b/config.example.py index c09548f..466edaf 100644 --- a/config.example.py +++ b/config.example.py @@ -14,4 +14,6 @@ 'botsfordiscord.com': None, }, }, + + 'user_agent': 'EmojiManagerBot (https://github.com/bmintz/emoji-manager-bot)', } diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..bfc85e1 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 +# encoding: utf-8 + +from .misc import * +from . import emote +from . import errors diff --git a/utils/emote.py b/utils/emote.py new file mode 100644 index 0000000..66fc02e --- /dev/null +++ b/utils/emote.py @@ -0,0 +1,19 @@ +#!/usr/bin/env python3 +# encoding: utf-8 + +import re + +""" +various utilities related to custom emotes +""" + +"""Matches :foo: and ;foo; but not :foo;. Used for emotes in text.""" +RE_EMOTE = re.compile(r'(:|;)(?P\w{2,32})\1|(?P\n)', re.ASCII) + +"""Matches only custom server emotes.""" +RE_CUSTOM_EMOTE = re.compile(r'<(?Pa?):(?P\w{2,32}):(?P\d{17,})>', re.ASCII) + +def url(id, *, animated: bool = False): + """Convert an emote ID to the image URL for that emote.""" + extension = 'gif' if animated else 'png' + return f'https://cdn.discordapp.com/emojis/{id}.{extension}?v=1' diff --git a/utils/errors.py b/utils/errors.py new file mode 100644 index 0000000..1fa8efb --- /dev/null +++ b/utils/errors.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +# encoding: utf-8 + +from discord.ext.commands import CommandError + + +class EmojiManagerError(CommandError): + """Generic error with the bot. This can be used to catch all bot errors.""" + pass + + +class HTTPException(EmojiManagerError): + """The server did not respond with an OK status code.""" + def __init__(self, status): + super().__init__(f'URL error: server returned error code {status}') + + +class EmojiNotFoundError(EmojiManagerError): + """An emoji with that name was not found""" + def __init__(self, name): + super().__init__(f'An emoji called `{name}` does not exist in this server.') + + +class InvalidImageError(EmojiManagerError): + """The image is not a GIF, PNG, or JPG""" + def __init__(self): + super().__init__('The image supplied was not a GIF, PNG, or JPG.') + + +class NoMoreSlotsError(EmojiManagerError): + """Raised in case all slots of a particular type (static/animated) are full""" + def __init__(self): + super().__init__('No more backend slots available.') + + +class PermissionDeniedError(EmojiManagerError): + """Raised when a user tries to modify an emoji without the Manage Emojis permission""" + def __init__(self, name): + super().__init__(f"You're not authorized to modify `{name}`.") + + +class DiscordError(Exception): + """Usually raised when the client cache is being baka""" + def __init__(self): + super().__init__('Discord seems to be having issues right now, please try again later.') diff --git a/utils/misc.py b/utils/misc.py new file mode 100644 index 0000000..b2a27a0 --- /dev/null +++ b/utils/misc.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +# encoding: utf-8 + +import discord + +"""various utilities for use within the bot""" + +"""Emotes used to indicate success/failure. You can obtain these from the discordbots.org guild, +but I uploaded them to my test server +so that both the staging and the stable versions of the bot can use them""" +SUCCESS_EMOTES = ('<:error:416845770239508512>', '<:success:416845760810844160>') + +def format_user(bot, id, *, mention=False): + """Format a user ID for human readable display.""" + user = bot.get_user(id) + if user is None: + return f'Unknown user with ID {id}' + # not mention: @null byte#8191 (140516693242937345) + # mention: <@140516693242937345> (null byte#8191) + # this allows people to still see the username and discrim + # if they don't share a server with that user + if mention: + return f'{user.mention} (@{user})' + else: + return f'@{user} ({user.id})' + +def format_http_exception(exception: discord.HTTPException): + """Formats a discord.HTTPException for relaying to the user. + Sample return value: + + BAD REQUEST (status code: 400): + Invalid Form Body + In image: File cannot be larger than 256 kb. + """ + return ( + f'{exception.response.reason} (status code: {exception.response.status}):' + f'\n{exception.text}') + +def strip_angle_brackets(string): + """Strip leading < and trailing > from a string. + Useful if a user sends you a url like to avoid embeds, or to convert emotes to reactions.""" + if string.startswith('<') and string.endswith('>'): + return string[1:-1] + return string