diff --git a/cogs/emote.py b/cogs/emote.py index abbb43b..e0c847c 100644 --- a/cogs/emote.py +++ b/cogs/emote.py @@ -37,8 +37,9 @@ from discord.ext import commands import utils import utils.image from utils import errors -from utils.converter import emote_type_filter_default from utils.paginator import ListPaginator +from utils.emote_client import EmoteClient +from utils.converter import emote_type_filter_default logger = logging.getLogger(__name__) @@ -75,6 +76,8 @@ class Emotes(commands.Cog): + self.bot.http.user_agent }) + self.emote_client = EmoteClient(token=self.bot.config['tokens']['discord']) + with open('data/ec-emotes-final.json') as f: self.ec_emotes = json.load(f) @@ -84,6 +87,7 @@ class Emotes(commands.Cog): def cog_unload(self): async def close(): await self.http.close() + await self.emote_client.close() for paginator in self.paginators: await paginator.stop() @@ -326,6 +330,8 @@ class Emotes(commands.Cog): if not url and not context.message.attachments: raise commands.BadArgument('A URL or attachment must be given.') + self.emote_client.check_create(context.guild.id) + url = url or context.message.attachments[0].url async with context.typing(): archive = await self.fetch_safe(url, valid_mimetypes=self.ARCHIVE_MIMETYPES) @@ -363,12 +369,13 @@ class Emotes(commands.Cog): async def add_safe(self, context, name, url, author_id, *, reason=None): """Try to add an emote. Returns a string that should be sent to the user.""" + self.emote_client.check_create(context.guild.id) try: image_data = await self.fetch_safe(url) except errors.InvalidFileError: raise errors.InvalidImageError - if type(image_data) is str: # error case + if type(image_data) is str: # error case (shitty i know) return image_data return await self.add_safe_bytes(context, name, author_id, image_data, reason=reason) @@ -436,8 +443,8 @@ class Emotes(commands.Cog): async def create_emote_from_bytes(self, guild, name, author_id, image_data: bytes, *, reason=None): image_data = await utils.image.resize_in_subprocess(image_data) if reason is None: - reason = f'Created by {utils.format_user(self.bot, author_id)}' - return await guild.create_custom_emoji(name=name, image=image_data, reason=reason) + reason = 'Created by ' + utils.format_user(self.bot, author_id) + return await self.emote_client.create(guild_id=guild.id, name=name, image=image_data, reason=reason) @commands.command(aliases=('delete', 'delet', 'rm')) async def remove(self, context, emote, *emotes): @@ -447,7 +454,11 @@ class Emotes(commands.Cog): """ if not emotes: emote = await self.parse_emote(context, emote) - await emote.delete(reason=f'Removed by {utils.format_user(self.bot, context.author.id)}') + await self.emote_client.delete( + guild_id=context.guild.id, + emote_id=emote.id, + reason='Removed by ' + utils.format_user(self.bot, context.author.id), + ) await context.send(fr'Emote \:{emote.name}: successfully removed.') else: for emote in (emote,) + emotes: diff --git a/utils/__init__.py b/utils/__init__.py index d26896f..f6d830d 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -5,5 +5,6 @@ from . import archive from . import emote from . import errors from . import paginator +from . import emote_client # note: do not import .image in case the user doesn't want it # since importing image can take a long time. diff --git a/utils/emote_client.py b/utils/emote_client.py new file mode 100644 index 0000000..dff1383 --- /dev/null +++ b/utils/emote_client.py @@ -0,0 +1,156 @@ +# © lambda#0987 +# SPDX-License-Identifier: AGPL-3.0-or-later + +import sys +import json +import asyncio +import aiohttp +import platform +import datetime +import urllib.parse +from typing import Dict +from http import HTTPStatus +import utils.image as image_utils +from utils.errors import RateLimitedError +from discord import HTTPException, Forbidden, NotFound, DiscordServerError + +class GuildRetryTimes: + """Holds the times, for a particular guild, + that we have to wait until for the rate limit for a particular HTTP method to elapse. + """ + __slots__ = frozenset({'POST', 'DELETE'}) + + def __init__(self, POST=None, DELETE=None): + self.POST = POST + self.DELETE = DELETE + + def validate(self): + now = datetime.datetime.now(tz=datetime.timezone.utc).timestamp() + if self.POST and self.POST < now: + self.POST = None + if self.DELETE and self.DELETE < now: + self.DELETE = None + return self.POST or self.DELETE + +GuildId = int + +async def json_or_text(resp): + text = await resp.text(encoding='utf-8') + try: + if resp.headers['content-type'] == 'application/json': + return json.loads(text) + except KeyError: + # Thanks Cloudflare + pass + + return text + +class EmoteClient: + BASE_URL = 'https://discord.com/api/v7' + HTTP_ERROR_CLASSES = { + HTTPStatus.FORBIDDEN: Forbidden, + HTTPStatus.NOT_FOUND: NotFound, + HTTPStatus.SERVICE_UNAVAILABLE: DiscordServerError, + } + + def __init__(self, *, token): + self.guild_rls: Dict[GuildId, GuildRetryTimes] = {} + user_agent = ( + 'EmoteManager-EmoteClient; ' + f'aiohttp/{aiohttp.__version__}; ' + f'{platform.python_implementation()}/{".".join(map(str, sys.version_info))}' + ) + self.http = aiohttp.ClientSession(headers={ + 'User-Agent': user_agent, + 'Authorization': 'Bot ' + token, + 'X-Ratelimit-Precision': 'millisecond', + }) + + async def request(self, method, path, guild_id, **kwargs): + self._check_rl(method, guild_id) + print('post check rl') + + headers = {} + # Emote Manager shouldn't use walrus op until Debian adopts 3.8 :( + reason = kwargs.pop('reason', None) + if reason: + headers['X-Audit-Log-Reason'] = urllib.parse.quote(reason, safe='/ ') + kwargs['headers'] = headers + + # TODO handle OSError and 500/502, like dpy does + async with self.http.request(method, self.BASE_URL + path, **kwargs) as resp: + if resp.status == HTTPStatus.TOO_MANY_REQUESTS: + print('handling rl') + return await self._handle_rl(resp, method, path, guild_id, **kwargs) + + data = await json_or_text(resp) + if resp.status in range(200, 300): + return data + + error_cls = self.HTTP_ERROR_CLASSES.get(resp.status, HTTPException) + raise error_cls(resp, data) + + def _check_rl(self, method, guild_id): + try: + rls = self.guild_rls[guild_id] + except KeyError: + print('guild not found', repr(guild_id)) + return + + if not rls.validate(): + print('guild rls invalid') + del self.guild_rls[guild_id] + return + + retry_at = getattr(rls, method, None) + if retry_at: + print('retry later') + raise RateLimitedError(retry_at) + + async def _handle_rl(self, resp, method, path, guild_id, **kwargs): + retry_after = (await resp.json())['retry_after'] / 1000.0 + retry_at = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=retry_after) + + # cache unconditionally in case request() is called again while we're sleeping + try: + rls = self.guild_rls[guild_id] + except KeyError: + self.guild_rls[guild_id] = rls = GuildRetryTimes() + + setattr(rls, resp.method, retry_at.timestamp()) + + if resp.method not in GuildRetryTimes.__slots__ or retry_after < 10.0: + await asyncio.sleep(retry_after) + # woo mutual recursion + return await self.request(method, path, guild_id, **kwargs) + + # we've been hit with one of those crazy high rate limits, which only occur for specific methods + raise RateLimitedError(retry_at) + + # optimization methods that let us check the RLs before downloading the user's image + def check_create(self, guild_id): + self._check_rl('POST', guild_id) + + def check_delete(self, guild_id): + self._check_rl('DELETE', guild_id) + + async def create(self, *, guild_id, name, image: bytes, role_ids=(), reason=None): + return await self.request( + 'POST', f'/guilds/{guild_id}/emojis', + guild_id, + json=dict(name=name, image=image_utils.image_to_base64_url(image), roles=role_ids), + reason=reason, + ) + + async def delete(self, *, guild_id, emote_id, reason=None): + return await self.request('DELETE', f'/guilds/{guild_id}/emojis/{emote_id}', guild_id, reason=reason) + + async def __aenter__(self): + self.http = await self.http.__aenter__() + return self + + async def __aexit__(self, *excinfo): + return await self.http.__aexit__(*excinfo) + + async def close(self): + return await self.http.close() diff --git a/utils/errors.py b/utils/errors.py index 67bdcdf..1e636c4 100644 --- a/utils/errors.py +++ b/utils/errors.py @@ -13,11 +13,11 @@ # You should have received a copy of the GNU Affero General Public License # along with Emote Manager. If not, see . -import asyncio - -from discord.ext import commands - import utils +import asyncio +import humanize +import datetime +from discord.ext import commands class MissingManageEmojisPermission(commands.MissingPermissions): """The invoker or the bot doesn't have permissions to manage server emojis.""" @@ -45,10 +45,19 @@ class ImageConversionTimeoutError(ImageProcessingTimeoutError): super().__init__('Error: converting the image to a GIF took too long.') class HTTPException(EmoteManagerError): - """The server did not respond with an OK status code.""" + """The server did not respond with an OK status code. This is only for non-Discord HTTP requests.""" def __init__(self, status): super().__init__(f'URL error: server returned error code {status}') +class RateLimitedError(EmoteManagerError): + def __init__(self, retry_at): + if isinstance(retry_at, float): + # it took me about an HOUR to realize i had to pass tz because utcfromtimestamp returns a NAÏVE time obj! + retry_at = datetime.datetime.fromtimestamp(retry_at, tz=datetime.timezone.utc) + # humanize.naturaltime is annoying to work with due to timezones so we use this + delta = humanize.naturaldelta(retry_at, when=datetime.datetime.now(tz=datetime.timezone.utc)) + super().__init__(f'Error: Discord told me to slow down! Please retry this command in {delta}.') + class EmoteNotFoundError(EmoteManagerError): """An emote with that name was not found""" def __init__(self, name):