From ddf325c57fba8b76f3067856153e70f3ac1547d5 Mon Sep 17 00:00:00 2001 From: bobloy Date: Tue, 14 Nov 2017 14:10:41 -0500 Subject: [PATCH] initial commit --- reactrestrict/__init__.py | 5 + reactrestrict/reactrestrict.py | 360 +++++++++++++++++++++++++++++++++ 2 files changed, 365 insertions(+) create mode 100644 reactrestrict/__init__.py create mode 100644 reactrestrict/reactrestrict.py diff --git a/reactrestrict/__init__.py b/reactrestrict/__init__.py new file mode 100644 index 0000000..8bd425f --- /dev/null +++ b/reactrestrict/__init__.py @@ -0,0 +1,5 @@ +from .reactrestrict import ReactRestrict + + +def setup(bot): + bot.add_cog(ReactRestrict(bot)) diff --git a/reactrestrict/reactrestrict.py b/reactrestrict/reactrestrict.py new file mode 100644 index 0000000..8d00336 --- /dev/null +++ b/reactrestrict/reactrestrict.py @@ -0,0 +1,360 @@ +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: + """ + This cog enables role assignment/removal based on reactions 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(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 + + return None + + # async def _wait_for_emoji(self, ctx: commands.Context): + # """ + # Asks the user to react to this message and returns the emoji string if unicode + # or ID if custom. + + # :param ctx: + # :raises asyncio.TimeoutError: + # If the user does not respond in time. + # :return: + # """ + # message = await ctx.send("Please react to this message with the reaction you" + # " would like to add/remove, you have 20 seconds to" + # " respond.") + + # def _wait_check(react, user): + # msg = react.message + # return msg.id == message.id and user.id == ctx.author.id + + # reaction, _ = await ctx.bot.wait_for('reaction_add', check=_wait_check, timeout=20) + + # try: + # ret = reaction.emoji.id + # except AttributeError: + The emoji is unicode + # ret = reaction.emoji + + # return ret, reaction.emoji + + @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.bot.send_cmd_help(ctx) + + @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.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) + + 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 apprroles in roles: + if apprrole in member.roles: + return + + message = await self._get_message(ctx, 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