EmoteManager/utils/emote_client.py

# © 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
from nextcord import PartialEmoji
import utils.image as image_utils
from utils.errors import RateLimitedError
from nextcord import HTTPException, Forbidden, NotFound, DiscordServerError

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/v9'
    HTTP_ERROR_CLASSES = {
        HTTPStatus.FORBIDDEN: Forbidden,
        HTTPStatus.NOT_FOUND: NotFound,
        HTTPStatus.SERVICE_UNAVAILABLE: DiscordServerError,
    }

    def __init__(self, bot):
        self.guild_rls: Dict[GuildId, float] = {}
        self.http = aiohttp.ClientSession(headers={
            'User-Agent': bot.config['user_agent'] + ' ' + bot.http.user_agent,
            'Authorization': 'Bot ' + bot.config['tokens']['discord'],
            'X-Ratelimit-Precision': 'millisecond',
        })

    async def request(self, method, path, guild_id, **kwargs):
        self.check_rl(guild_id)

        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:
                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)

    # optimization method that lets us check the RL before downloading the user's image.
    # also lets us preemptively check the RL before doing a request
    def check_rl(self, guild_id):
        try:
            retry_at = self.guild_rls[guild_id]
        except KeyError:
            return

        now = datetime.datetime.now(tz=datetime.timezone.utc).timestamp()
        if retry_at < now:
            del self.guild_rls[guild_id]
            return

        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
        self.guild_rls[guild_id] = retry_at.timestamp()

        if 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)

    async def create(self, *, guild, name, image: bytes, role_ids=(), reason=None):
        data = 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,
        )
        return PartialEmoji(animated=data.get('animated', False), name=data.get('name'), id=data.get('id'))

    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()