Fox-V3/reactrestrict/reactrestrict.py

import asyncio
from typing import List, Union

import discord
from discord.ext import commands

from redbot.core import Config
from redbot.core.bot import Red


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:
    """
    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.

        :param int message_id:
        :param str or int emoji:
        :param discord.Role role:
        """
        # 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 int message_id:
        :param str or int emoji:
        :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 int message_id:
        :param str or int emoji:
        :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.

        :param ctx:
        :param message_id:
        :return:
        """
        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:
            await ctx.send_help()

    @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