diff options
-rw-r--r-- | bot/__main__.py | 1 | ||||
-rw-r--r-- | bot/cogs/duck_pond.py | 206 |
2 files changed, 207 insertions, 0 deletions
diff --git a/bot/__main__.py b/bot/__main__.py index f352cd60e..ea7c43a12 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -55,6 +55,7 @@ if not DEBUG_MODE: bot.load_extension("bot.cogs.alias") bot.load_extension("bot.cogs.defcon") bot.load_extension("bot.cogs.eval") +bot.load_extension("bot.cogs.duck_pond") bot.load_extension("bot.cogs.free") bot.load_extension("bot.cogs.information") bot.load_extension("bot.cogs.jams") diff --git a/bot/cogs/duck_pond.py b/bot/cogs/duck_pond.py new file mode 100644 index 000000000..d5d528458 --- /dev/null +++ b/bot/cogs/duck_pond.py @@ -0,0 +1,206 @@ +import logging +from typing import List, Optional, Union + +import discord +from discord import Color, Embed, Member, Message, PartialEmoji, RawReactionActionEvent, Reaction, User, errors +from discord.ext.commands import Bot, Cog + +import bot.constants as constants +from bot.utils.messages import send_attachments + +log = logging.getLogger(__name__) + + +class DuckPond(Cog): + """Relays messages to #duck-pond whenever a certain number of duck reactions have been achieved.""" + + def __init__(self, bot: Bot): + self.bot = bot + self.log = log + self.webhook_id = constants.Webhooks.duck_pond + self.bot.loop.create_task(self.fetch_webhook()) + + async def fetch_webhook(self): + """Fetches the webhook object, so we can post to it.""" + await self.bot.wait_until_ready() + + try: + self.webhook = await self.bot.fetch_webhook(self.webhook_id) + except discord.HTTPException: + self.log.exception(f"Failed to fetch webhook with id `{self.webhook_id}`") + + @staticmethod + def is_staff(member: Union[User, Member]) -> bool: + """Check if a specific member or user is staff""" + if hasattr(member, "roles"): + for role in member.roles: + if role.id in constants.STAFF_ROLES: + return True + return False + + @staticmethod + def has_green_checkmark(message: Optional[Message] = None, reaction_list: Optional[List[Reaction]] = None) -> bool: + """Check if the message has a green checkmark reaction.""" + assert message or reaction_list, "You can either pass message or reactions, but not both, or neither." + + if message: + reactions = message.reactions + else: + reactions = reaction_list + + for reaction in reactions: + if isinstance(reaction.emoji, str): + if reaction.emoji == "✅": + return True + elif isinstance(reaction.emoji, PartialEmoji): + if reaction.emoji.name == "✅": + return True + return False + + async def send_webhook( + self, + content: Optional[str] = None, + username: Optional[str] = None, + avatar_url: Optional[str] = None, + embed: Optional[Embed] = None, + ) -> None: + try: + await self.webhook.send( + content=content, + username=username, + avatar_url=avatar_url, + embed=embed + ) + except discord.HTTPException as exc: + self.log.exception( + f"Failed to send a message to the Duck Pool webhook", + exc_info=exc + ) + + async def count_ducks(self, message: Optional[Message] = None, reaction_list: Optional[List[Reaction]] = None) -> int: + """Count the number of ducks in the reactions of a specific message. + + Only counts ducks added by staff members. + """ + assert message or reaction_list, "You can either pass message or reactions, but not both, or neither." + + duck_count = 0 + duck_reactors = [] + + if message: + reactions = message.reactions + else: + reactions = reaction_list + + for reaction in reactions: + async for user in reaction.users(): + + # Is the user or member a staff member? + if self.is_staff(user) and user.id not in duck_reactors: + + # Is the emoji a duck? + if hasattr(reaction.emoji, "id"): + if reaction.emoji.id in constants.DuckPond.duck_custom_emojis: + duck_count += 1 + duck_reactors.append(user.id) + else: + if isinstance(reaction.emoji, str): + if reaction.emoji == "🦆": + duck_count += 1 + duck_reactors.append(user.id) + elif isinstance(reaction.emoji, PartialEmoji): + if reaction.emoji.name == "🦆": + duck_count += 1 + duck_reactors.append(user.id) + return duck_count + + @Cog.listener() + async def on_raw_reaction_add(self, payload: RawReactionActionEvent) -> None: + """Determine if a message should be sent to the duck pond. + + This will count the number of duck reactions on the message, and if this amount meets the + amount of ducks specified in the config under duck_pond/ducks_required, it will + send the message off to the duck pond. + """ + channel = discord.utils.get(self.bot.get_all_channels(), id=payload.channel_id) + message = await channel.fetch_message(payload.message_id) + member = discord.utils.get(message.guild.members, id=payload.user_id) + + # Is the member a staff member? + if not self.is_staff(member): + return + + # Bot reactions don't count. + if member.bot: + return + + # Is the emoji in the reaction a duck? + if payload.emoji.is_custom_emoji(): + if payload.emoji.id not in constants.DuckPond.duck_custom_emojis: + return + else: + if payload.emoji.name != "🦆": + return + + # Does the message already have a green checkmark? + if self.has_green_checkmark(message): + return + + # Time to count our ducks! + duck_count = await self.count_ducks(message) + + # If we've got more than the required amount of ducks, send the message to the duck_pond. + if duck_count >= constants.DuckPond.ducks_required: + clean_content = message.clean_content + + if clean_content: + await self.send_webhook( + content=message.clean_content, + username=message.author.display_name, + avatar_url=message.author.avatar_url + ) + + if message.attachments: + try: + await send_attachments(message, self.webhook) + except (errors.Forbidden, errors.NotFound): + e = Embed( + description=":x: **This message contained an attachment, but it could not be retrieved**", + color=Color.red() + ) + await self.send_webhook( + embed=e, + username=message.author.display_name, + avatar_url=message.author.avatar_url + ) + except discord.HTTPException as exc: + self.log.exception( + f"Failed to send an attachment to the webhook", + exc_info=exc + ) + await message.add_reaction("✅") + + @Cog.listener() + async def on_raw_reaction_remove(self, payload: RawReactionActionEvent) -> None: + """Ensure that people don't remove the green checkmark from duck ponded messages.""" + channel = discord.utils.get(self.bot.get_all_channels(), id=payload.channel_id) + message = await channel.fetch_message(payload.message_id) + + # Prevent the green checkmark from being removed + if isinstance(payload.emoji, str): + if payload.emoji == "✅": + duck_count = await self.count_ducks(message) + if duck_count >= constants.DuckPond.ducks_required: + await message.add_reaction("✅") + + elif isinstance(payload.emoji, PartialEmoji): + if payload.emoji.name == "✅": + duck_count = await self.count_ducks(message) + if duck_count >= constants.DuckPond.ducks_required: + await message.add_reaction("✅") + + +def setup(bot: Bot) -> None: + """Token Remover cog load.""" + bot.add_cog(DuckPond(bot)) + log.info("Cog loaded: DuckPond") |