Fox-V3/reactrestrict/reactrestrict.py

from typing import List, Union

import discord
from redbot.core import Config
from redbot.core import commands
from redbot.core.bot import Red
from typing import Any

Cog: Any = getattr(commands, "Cog", object)


class ReactRestrictCombo:
    def __init__(self, message_id, role_id):
        self.message_id = message_id
        self.role_id = role_id

    def __eq__(self, other: "ReactRestrictCombo"):
        return self.message_id == other.message_id and self.role_id == other.role_id

    def to_json(self):
        return {"message_id": self.message_id, "role_id": self.role_id}

    @classmethod
    def from_json(cls, data):
        return cls(data["message_id"], data["role_id"])


class ReactRestrict(Cog):
    """
    Prevent specific roles from reacting to specific messages
    """

    def __init__(self, red: Red):
        self.bot = red
        self.config = Config.get_conf(
            self, 8210197991168210111511611410599116, force_registration=True
        )
        self.config.register_global(registered_combos=[])

    async def combo_list(self) -> List[ReactRestrictCombo]:
        """
        Returns a list of reactrestrict combos.

        :return:
        """
        cmd = self.config.registered_combos()

        return [ReactRestrictCombo.from_json(data) for data in await cmd]

    async def set_combo_list(self, combo_list: List[ReactRestrictCombo]):
        """
        Helper method to set the list of reactrestrict combos.

        :param combo_list:
        :return:
        """
        raw = [combo.to_json() for combo in combo_list]
        await self.config.registered_combos.set(raw)

    async def is_registered(self, message_id: int) -> bool:
        """
        Determines if a message ID has been registered.

        :param message_id:
        :return:
        """
        return any(message_id == combo.message_id for combo in await self.combo_list())

    async def add_reactrestrict(self, message_id: int, role: discord.Role):
        """
        Adds a react|role combo.
        """
        # is_custom = True
        # if isinstance(emoji, str):
        # is_custom = False

        combo = ReactRestrictCombo(message_id, role.id)

        current_combos = await self.combo_list()

        if combo not in current_combos:
            current_combos.append(combo)
            await self.set_combo_list(current_combos)

    async def remove_react(self, message_id: int, role: discord.Role):
        """
        Removes a given reaction

        :param message_id:
        :param role:
        :return:
        """
        current_combos = await self.combo_list()

        to_keep = [
            c for c in current_combos if not (c.message_id == message_id and c.role_id == role.id)
        ]

        if to_keep != current_combos:
            await self.set_combo_list(to_keep)

    async def has_reactrestrict_combo(self, message_id: int) -> (bool, List[ReactRestrictCombo]):
        """
         Determines if there is an existing role combo for a given message
        and emoji ID.

        :param message_id:
        :return:
        """
        if not await self.is_registered(message_id):
            return False, []

        combos = await self.combo_list()

        ret = [c for c in combos if c.message_id == message_id]

        return len(ret) > 0, ret

    def _get_member(self, channel_id: int, user_id: int) -> discord.Member:
        """
        Tries to get a member with the given user ID from the guild that has
        the given channel ID.

        :param int channel_id:
        :param int user_id:
        :rtype:
            discord.Member
        :raises LookupError:
            If no such channel or member can be found.
        """
        channel = self.bot.get_channel(channel_id)
        try:
            member = channel.guild.get_member(user_id)
        except AttributeError as e:
            raise LookupError("No channel found.") from e

        if member is None:
            raise LookupError("No member found.")

        return member

    def _get_role(self, guild: discord.Guild, role_id: int) -> discord.Role:
        """
        Gets a role object from the given guild with the given ID.

        :param discord.Guild guild:
        :param int role_id:
        :rtype:
            discord.Role
        :raises LookupError:
            If no such role exists.
        """
        role = discord.utils.get(guild.roles, id=role_id)

        if role is None:
            raise LookupError("No role found.")

        return role

    async def _get_message_from_channel(
        self, channel_id: int, message_id: int
    ) -> Union[discord.Message, None]:
        """
        Tries to find a message by ID in the current guild context.
        """
        channel = self.bot.get_channel(channel_id)
        try:
            return await channel.get_message(message_id)
        except discord.NotFound:
            pass
        except AttributeError:  # VoiceChannel object has no attribute 'get_message'
            pass

        return None

    async def _get_message(
        self, ctx: commands.Context, message_id: int
    ) -> Union[discord.Message, None]:
        """
        Tries to find a message by ID in the current guild context.

        :param ctx:
        :param message_id:
        :return:
        """
        for channel in ctx.guild.channels:
            try:
                return await channel.get_message(message_id)
            except discord.NotFound:
                pass
            except AttributeError:  # VoiceChannel object has no attribute 'get_message'
                pass
            except discord.Forbidden:  # No access to channel, skip
                pass

        return None

    @commands.group()
    async def reactrestrict(self, ctx: commands.Context):
        """
        Base command for this cog. Check help for the commands list.
        """
        if ctx.invoked_subcommand is None:
            pass

    @reactrestrict.command()
    async def add(self, ctx: commands.Context, message_id: int, *, role: discord.Role):
        """
        Adds a reaction|role combination to a registered message, don't use
        quotes for the role name.
        """
        message = await self._get_message(ctx, message_id)
        if message is None:
            await ctx.send("That message doesn't seem to exist.")
            return

        # try:
        #     emoji, actual_emoji = await self._wait_for_emoji(ctx)
        # except asyncio.TimeoutError:
        #     await ctx.send("You didn't respond in time, please redo this command.")
        #     return
        #
        # try:
        #     await message.add_reaction(actual_emoji)
        # except discord.HTTPException:
        #     await ctx.send("I can't add that emoji because I'm not in the guild that"
        #                    " owns it.")
        #     return
        #
        # noinspection PyTypeChecker
        await self.add_reactrestrict(message_id, role)

        await ctx.send("Message|Role combo added.")

    @reactrestrict.command()
    async def remove(self, ctx: commands.Context, message_id: int, role: discord.Role):
        """
        Removes role associated with a given reaction.
        """
        # try:
        #     emoji, actual_emoji = await self._wait_for_emoji(ctx)
        # except asyncio.TimeoutError:
        #     await ctx.send("You didn't respond in time, please redo this command.")
        #     return

        # noinspection PyTypeChecker
        await self.remove_react(message_id, role)

        await ctx.send("Reaction removed.")

    async def on_raw_reaction_add(
        self, emoji: discord.PartialEmoji, message_id: int, channel_id: int, user_id: int
    ):
        """
        Event handler for long term reaction watching.

        :param discord.PartialReactionEmoji emoji:
        :param int message_id:
        :param int channel_id:
        :param int user_id:
        :return:
        """
        if emoji.is_custom_emoji():
            emoji_id = emoji.id
        else:
            emoji_id = emoji.name

        has_reactrestrict, combos = await self.has_reactrestrict_combo(message_id)

        if not has_reactrestrict:
            return

        try:
            member = self._get_member(channel_id, user_id)
        except LookupError:
            return

        if member.bot:
            return

        try:
            roles = [self._get_role(member.guild, c.role_id) for c in combos]
        except LookupError:
            return

        for apprrole in roles:
            if apprrole in member.roles:
                return

        message = await self._get_message_from_channel(channel_id, message_id)
        await message.remove_reaction(emoji, member)

    #     try:
    #         await member.add_roles(*roles)
    #     except discord.Forbidden:
    #         pass
    #
    # async def on_raw_reaction_remove(self, emoji: discord.PartialReactionEmoji,
    #                                  message_id: int, channel_id: int, user_id: int):
    #     """
    #     Event handler for long term reaction watching.
    #
    #     :param discord.PartialReactionEmoji emoji:
    #     :param int message_id:
    #     :param int channel_id:
    #     :param int user_id:
    #     :return:
    #     """
    #     if emoji.is_custom_emoji():
    #         emoji_id = emoji.id
    #     else:
    #         emoji_id = emoji.name
    #
    #     has_reactrestrict, combos = await self.has_reactrestrict_combo(message_id, emoji_id)
    #
    #     if not has_reactrestrict:
    #         return
    #
    #     try:
    #         member = self._get_member(channel_id, user_id)
    #     except LookupError:
    #         return
    #
    #     if member.bot:
    #         return
    #
    #     try:
    #         roles = [self._get_role(member.guild, c.role_id) for c in combos]
    #     except LookupError:
    #         return
    #
    #     try:
    #         await member.remove_roles(*roles)
    #     except discord.Forbidden:
    #         pass