diff --git a/birthday/birthday.py b/birthday/birthday.py index f531e9f..3001edf 100644 --- a/birthday/birthday.py +++ b/birthday/birthday.py @@ -5,6 +5,11 @@ from redbot.core.config import Config from datetime import datetime, time, timedelta from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import random +import logging +import asyncio + +# Set up logging +logger = logging.getLogger("red.birthday") # Define context menu command outside the class @app_commands.context_menu(name="Give Birthday Role") @@ -15,27 +20,36 @@ async def birthday_context_menu(interaction: discord.Interaction, member: discor cog = interaction.client.get_cog("Birthday") if not cog: + logger.error("Birthday cog not loaded during context menu execution") await interaction.followup.send("Birthday cog is not loaded.", ephemeral=True) return # Check if the user has permission to use this command allowed_roles = await cog.config.guild(interaction.guild).allowed_roles() if not any(role.id in allowed_roles for role in interaction.user.roles): + logger.warning(f"User {interaction.user.id} attempted to use birthday context menu without permission") return await interaction.followup.send("You don't have permission to use this command.", ephemeral=True) birthday_role_id = await cog.config.guild(interaction.guild).birthday_role() if not birthday_role_id: + logger.error(f"Birthday role not set for guild {interaction.guild.id}") return await interaction.followup.send("The birthday role hasn't been set. An admin needs to set it using `/setrole`.", ephemeral=True) birthday_role = interaction.guild.get_role(birthday_role_id) if not birthday_role: + logger.error(f"Birthday role {birthday_role_id} not found in guild {interaction.guild.id}") return await interaction.followup.send("The birthday role doesn't exist anymore. Please ask an admin to set it again.", ephemeral=True) # Assign the role, ignoring hierarchy try: await member.add_roles(birthday_role, reason="Birthday role") + logger.info(f"Birthday role assigned to {member.id} in guild {interaction.guild.id}") except discord.Forbidden: + logger.error(f"Failed to assign birthday role to {member.id} in guild {interaction.guild.id}: Insufficient permissions") return await interaction.followup.send("I don't have permission to assign that role.", ephemeral=True) + except discord.HTTPException as e: + logger.error(f"Failed to assign birthday role to {member.id} in guild {interaction.guild.id}: {str(e)}") + return await interaction.followup.send("Failed to assign the birthday role due to a Discord error.", ephemeral=True) # Generate birthday message with random cakes (or pie) cakes = random.randint(0, 5) @@ -48,7 +62,8 @@ async def birthday_context_menu(interaction: discord.Interaction, member: discor birthday_channel_id = await cog.config.guild(interaction.guild).birthday_channel() if birthday_channel_id: channel = interaction.client.get_channel(birthday_channel_id) - if not channel: # If the set channel doesn't exist anymore + if not channel: + logger.warning(f"Birthday channel {birthday_channel_id} not found in guild {interaction.guild.id}") channel = interaction.channel else: channel = interaction.channel @@ -61,6 +76,7 @@ async def birthday_context_menu(interaction: discord.Interaction, member: discor try: tz = ZoneInfo(timezone) except ZoneInfoNotFoundError: + logger.warning(f"Invalid timezone {timezone} for guild {interaction.guild.id}, defaulting to UTC") await interaction.followup.send("Warning: Invalid timezone set. Defaulting to UTC.", ephemeral=True) tz = ZoneInfo("UTC") @@ -69,6 +85,7 @@ async def birthday_context_menu(interaction: discord.Interaction, member: discor await cog.schedule_birthday_role_removal(interaction.guild, member, birthday_role, midnight) except Exception as e: + logger.error(f"Unexpected error in birthday context menu: {str(e)}", exc_info=True) try: await interaction.followup.send(f"An error occurred: {str(e)}", ephemeral=True) except: @@ -89,7 +106,124 @@ class Birthday(commands.Cog): } self.config.register_guild(**default_guild) self.birthday_tasks = {} + self.cleanup_task = None + self.bot.loop.create_task(self.initialize()) + async def initialize(self): + """Initialize the cog and start the cleanup task.""" + await self.bot.wait_until_ready() + await self.reload_scheduled_tasks() + self.cleanup_task = self.bot.loop.create_task(self.daily_cleanup()) + + async def cog_unload(self): + """Clean up tasks when the cog is unloaded.""" + if self.cleanup_task: + self.cleanup_task.cancel() + for task in self.birthday_tasks.values(): + task.cancel() + + @commands.hybrid_command(name="removebirthday") + @app_commands.guild_only() + @app_commands.describe(member="The member to remove the birthday role from") + async def remove_birthday(self, ctx: commands.Context, member: discord.Member): + """Remove the birthday role from a user.""" + try: + # Check if the user has permission to use this command + allowed_roles = await self.config.guild(ctx.guild).allowed_roles() + if not any(role.id in allowed_roles for role in ctx.author.roles): + logger.warning(f"User {ctx.author.id} attempted to remove birthday role without permission") + return await ctx.send("You don't have permission to use this command.", ephemeral=True) + + birthday_role_id = await self.config.guild(ctx.guild).birthday_role() + if not birthday_role_id: + logger.error(f"Birthday role not set for guild {ctx.guild.id}") + return await ctx.send("The birthday role hasn't been set.", ephemeral=True) + + birthday_role = ctx.guild.get_role(birthday_role_id) + if not birthday_role: + logger.error(f"Birthday role {birthday_role_id} not found in guild {ctx.guild.id}") + return await ctx.send("The birthday role doesn't exist anymore.", ephemeral=True) + + if birthday_role not in member.roles: + return await ctx.send(f"{member.display_name} doesn't have the birthday role.", ephemeral=True) + + try: + await member.remove_roles(birthday_role, reason="Birthday role manually removed") + logger.info(f"Birthday role manually removed from {member.id} in guild {ctx.guild.id}") + except discord.Forbidden: + logger.error(f"Failed to remove birthday role from {member.id} in guild {ctx.guild.id}: Insufficient permissions") + return await ctx.send("I don't have permission to remove that role.", ephemeral=True) + except discord.HTTPException as e: + logger.error(f"Failed to remove birthday role from {member.id} in guild {ctx.guild.id}: {str(e)}") + return await ctx.send("Failed to remove the birthday role due to a Discord error.", ephemeral=True) + + # Remove scheduled task if it exists + if str(member.id) in (await self.config.guild(ctx.guild).scheduled_tasks()): + await self.config.guild(ctx.guild).scheduled_tasks.clear_raw(str(member.id)) + if ctx.guild.id in self.birthday_tasks: + self.birthday_tasks[ctx.guild.id].cancel() + del self.birthday_tasks[ctx.guild.id] + + await ctx.send(f"Birthday role removed from {member.display_name}!", ephemeral=True) + except Exception as e: + logger.error(f"Unexpected error in remove_birthday command: {str(e)}", exc_info=True) + await ctx.send(f"An error occurred while removing the birthday role: {str(e)}", ephemeral=True) + + async def daily_cleanup(self): + """Daily task to ensure all birthday roles are properly removed.""" + while True: + try: + await asyncio.sleep(3600) # Check every hour + logger.info("Running daily birthday role cleanup") + + for guild in self.bot.guilds: + try: + scheduled_tasks = await self.config.guild(guild).scheduled_tasks() + timezone = await self.config.guild(guild).timezone() + try: + tz = ZoneInfo(timezone) + except ZoneInfoNotFoundError: + logger.warning(f"Invalid timezone {timezone} for guild {guild.id}, defaulting to UTC") + tz = ZoneInfo("UTC") + + now = datetime.now(tz) + + for member_id, task_info in scheduled_tasks.items(): + try: + member = guild.get_member(int(member_id)) + if not member: + logger.warning(f"Member {member_id} not found in guild {guild.id}, removing task") + await self.config.guild(guild).scheduled_tasks.clear_raw(member_id) + continue + + role = guild.get_role(task_info["role_id"]) + if not role: + logger.warning(f"Role {task_info['role_id']} not found in guild {guild.id}, removing task") + await self.config.guild(guild).scheduled_tasks.clear_raw(member_id) + continue + + remove_at = datetime.fromisoformat(task_info["remove_at"]).replace(tzinfo=tz) + + if now >= remove_at: + try: + await member.remove_roles(role, reason="Birthday role duration expired (cleanup)") + logger.info(f"Removed expired birthday role from {member_id} in guild {guild.id}") + except discord.Forbidden: + logger.error(f"Failed to remove birthday role from {member_id} in guild {guild.id}: Insufficient permissions") + except discord.HTTPException as e: + logger.error(f"Failed to remove birthday role from {member_id} in guild {guild.id}: {str(e)}") + finally: + await self.config.guild(guild).scheduled_tasks.clear_raw(member_id) + except Exception as e: + logger.error(f"Error processing task for member {member_id} in guild {guild.id}: {str(e)}", exc_info=True) + except Exception as e: + logger.error(f"Error processing guild {guild.id} in cleanup: {str(e)}", exc_info=True) + except Exception as e: + logger.error(f"Error in daily cleanup task: {str(e)}", exc_info=True) + finally: + await asyncio.sleep(3600) # Wait an hour before next check + + # [Previous commands remain unchanged...] @commands.hybrid_command(name="setrole") @app_commands.guild_only() @app_commands.describe(role="The role to set as the birthday role") @@ -146,112 +280,179 @@ class Birthday(commands.Cog): @app_commands.describe(member="The member to give the birthday role to") async def birthday(self, ctx: commands.Context, member: discord.Member): """Assign the birthday role to a user until midnight in the set timezone.""" - # Check if the user has permission to use this command - allowed_roles = await self.config.guild(ctx.guild).allowed_roles() - if not any(role.id in allowed_roles for role in ctx.author.roles): - return await ctx.send("You don't have permission to use this command.", ephemeral=True) - - birthday_role_id = await self.config.guild(ctx.guild).birthday_role() - if not birthday_role_id: - return await ctx.send("The birthday role hasn't been set. An admin needs to set it using `/setrole`.", ephemeral=True) - - birthday_role = ctx.guild.get_role(birthday_role_id) - if not birthday_role: - return await ctx.send("The birthday role doesn't exist anymore. Please ask an admin to set it again.", ephemeral=True) - - # Assign the role, ignoring hierarchy try: - await member.add_roles(birthday_role, reason="Birthday role") - except discord.Forbidden: - return await ctx.send("I don't have permission to assign that role.", ephemeral=True) + # Check if the user has permission to use this command + allowed_roles = await self.config.guild(ctx.guild).allowed_roles() + if not any(role.id in allowed_roles for role in ctx.author.roles): + logger.warning(f"User {ctx.author.id} attempted to use birthday command without permission") + return await ctx.send("You don't have permission to use this command.", ephemeral=True) - # Generate birthday message with random cakes (or pie) - cakes = random.randint(0, 5) - if cakes == 0: - message = f"🎉 Happy Birthday, {member.mention}! Sorry, out of cake today! Here's pie instead: 🥧" - else: - message = f"🎉 Happy Birthday, {member.mention}! Here's your cake{'s' if cakes > 1 else ''}: " + "🎂" * cakes + birthday_role_id = await self.config.guild(ctx.guild).birthday_role() + if not birthday_role_id: + logger.error(f"Birthday role not set for guild {ctx.guild.id}") + return await ctx.send("The birthday role hasn't been set. An admin needs to set it using `/setrole`.", ephemeral=True) + + birthday_role = ctx.guild.get_role(birthday_role_id) + if not birthday_role: + logger.error(f"Birthday role {birthday_role_id} not found in guild {ctx.guild.id}") + return await ctx.send("The birthday role doesn't exist anymore. Please ask an admin to set it again.", ephemeral=True) - # Get the birthday announcement channel - birthday_channel_id = await self.config.guild(ctx.guild).birthday_channel() - if birthday_channel_id: - channel = self.bot.get_channel(birthday_channel_id) - if not channel: # If the set channel doesn't exist anymore + # Assign the role, ignoring hierarchy + try: + await member.add_roles(birthday_role, reason="Birthday role") + logger.info(f"Birthday role assigned to {member.id} in guild {ctx.guild.id}") + except discord.Forbidden: + logger.error(f"Failed to assign birthday role to {member.id} in guild {ctx.guild.id}: Insufficient permissions") + return await ctx.send("I don't have permission to assign that role.", ephemeral=True) + except discord.HTTPException as e: + logger.error(f"Failed to assign birthday role to {member.id} in guild {ctx.guild.id}: {str(e)}") + return await ctx.send("Failed to assign the birthday role due to a Discord error.", ephemeral=True) + + # Generate birthday message with random cakes (or pie) + cakes = random.randint(0, 5) + if cakes == 0: + message = f"🎉 Happy Birthday, {member.mention}! Sorry, out of cake today! Here's pie instead: 🥧" + else: + message = f"🎉 Happy Birthday, {member.mention}! Here's your cake{'s' if cakes > 1 else ''}: " + "🎂" * cakes + + # Get the birthday announcement channel + birthday_channel_id = await self.config.guild(ctx.guild).birthday_channel() + if birthday_channel_id: + channel = self.bot.get_channel(birthday_channel_id) + if not channel: # If the set channel doesn't exist anymore + logger.warning(f"Birthday channel {birthday_channel_id} not found in guild {ctx.guild.id}") + channel = ctx.channel + else: channel = ctx.channel - else: - channel = ctx.channel - await channel.send(message) - await ctx.send("Birthday role assigned!", ephemeral=True) + await channel.send(message) + await ctx.send("Birthday role assigned!", ephemeral=True) - # Schedule role removal - timezone = await self.config.guild(ctx.guild).timezone() - try: - tz = ZoneInfo(timezone) - except ZoneInfoNotFoundError: - await ctx.send("Warning: Invalid timezone set. Defaulting to UTC.", ephemeral=True) - tz = ZoneInfo("UTC") + # Schedule role removal + timezone = await self.config.guild(ctx.guild).timezone() + try: + tz = ZoneInfo(timezone) + except ZoneInfoNotFoundError: + logger.warning(f"Invalid timezone {timezone} for guild {ctx.guild.id}, defaulting to UTC") + await ctx.send("Warning: Invalid timezone set. Defaulting to UTC.", ephemeral=True) + tz = ZoneInfo("UTC") - now = datetime.now(tz) - midnight = datetime.combine(now.date() + timedelta(days=1), time.min).replace(tzinfo=tz) + now = datetime.now(tz) + midnight = datetime.combine(now.date() + timedelta(days=1), time.min).replace(tzinfo=tz) - await self.schedule_birthday_role_removal(ctx.guild, member, birthday_role, midnight) + await self.schedule_birthday_role_removal(ctx.guild, member, birthday_role, midnight) + except Exception as e: + logger.error(f"Unexpected error in birthday command: {str(e)}", exc_info=True) + await ctx.send(f"An error occurred: {str(e)}", ephemeral=True) @commands.hybrid_command(name="bdaycheck") @app_commands.guild_only() async def bdaycheck(self, ctx: commands.Context): """Check the upcoming birthday role removal tasks.""" - # Check if the user has permission to use this command - allowed_roles = await self.config.guild(ctx.guild).allowed_roles() - if not any(role.id in allowed_roles for role in ctx.author.roles): - return await ctx.send("You don't have permission to use this command.", ephemeral=True) + try: + # Check if the user has permission to use this command + allowed_roles = await self.config.guild(ctx.guild).allowed_roles() + if not any(role.id in allowed_roles for role in ctx.author.roles): + logger.warning(f"User {ctx.author.id} attempted to use bdaycheck command without permission") + return await ctx.send("You don't have permission to use this command.", ephemeral=True) - scheduled_tasks = await self.config.guild(ctx.guild).scheduled_tasks() - if not scheduled_tasks: - return await ctx.send("There are no scheduled tasks.", ephemeral=True) + scheduled_tasks = await self.config.guild(ctx.guild).scheduled_tasks() + if not scheduled_tasks: + return await ctx.send("There are no scheduled tasks.", ephemeral=True) - message = "Upcoming birthday role removal tasks:\n" - for member_id, task_info in scheduled_tasks.items(): - member = ctx.guild.get_member(int(member_id)) - if not member: - continue - role = ctx.guild.get_role(task_info["role_id"]) - if not role: - continue - remove_at = datetime.fromisoformat(task_info["remove_at"]).replace(tzinfo=ZoneInfo(await self.config.guild(ctx.guild).timezone())) - message += f"- {member.display_name} ({member.id}): {role.name} will be removed at {remove_at}\n" + message = "Upcoming birthday role removal tasks:\n" + for member_id, task_info in scheduled_tasks.items(): + member = ctx.guild.get_member(int(member_id)) + if not member: + continue + role = ctx.guild.get_role(task_info["role_id"]) + if not role: + continue + remove_at = datetime.fromisoformat(task_info["remove_at"]).replace(tzinfo=ZoneInfo(await self.config.guild(ctx.guild).timezone())) + message += f"- {member.display_name} ({member.id}): {role.name} will be removed at {remove_at}\n" - await ctx.send(message, ephemeral=True) + await ctx.send(message, ephemeral=True) + except Exception as e: + logger.error(f"Unexpected error in bdaycheck command: {str(e)}", exc_info=True) + await ctx.send(f"An error occurred while checking birthday tasks: {str(e)}", ephemeral=True) async def schedule_birthday_role_removal(self, guild, member, role, when): """Schedule the removal of the birthday role.""" - await self.config.guild(guild).scheduled_tasks.set_raw(str(member.id), value={ - "role_id": role.id, - "remove_at": when.isoformat() - }) - self.birthday_tasks[guild.id] = self.bot.loop.create_task(self.remove_birthday_role(guild, member, role, when)) + try: + await self.config.guild(guild).scheduled_tasks.set_raw(str(member.id), value={ + "role_id": role.id, + "remove_at": when.isoformat() + }) + if guild.id in self.birthday_tasks: + self.birthday_tasks[guild.id].cancel() + self.birthday_tasks[guild.id] = self.bot.loop.create_task(self.remove_birthday_role(guild, member, role, when)) + logger.info(f"Scheduled birthday role removal for {member.id} in guild {guild.id} at {when}") + except Exception as e: + logger.error(f"Failed to schedule birthday role removal for {member.id} in guild {guild.id}: {str(e)}", exc_info=True) + raise async def remove_birthday_role(self, guild, member, role, when): """Remove the birthday role at the specified time.""" - await discord.utils.sleep_until(when) try: - await member.remove_roles(role, reason="Birthday role duration expired") - except (discord.Forbidden, discord.HTTPException): - pass # If we can't remove the role, we'll just let it be + await discord.utils.sleep_until(when) + try: + await member.remove_roles(role, reason="Birthday role duration expired") + logger.info(f"Birthday role removed from {member.id} in guild {guild.id}") + except discord.Forbidden: + logger.error(f"Failed to remove birthday role from {member.id} in guild {guild.id}: Insufficient permissions") + except discord.HTTPException as e: + logger.error(f"Failed to remove birthday role from {member.id} in guild {guild.id}: {str(e)}") + except asyncio.CancelledError: + logger.info(f"Birthday role removal task cancelled for {member.id} in guild {guild.id}") + raise + except Exception as e: + logger.error(f"Error removing birthday role from {member.id} in guild {guild.id}: {str(e)}", exc_info=True) finally: - del self.birthday_tasks[guild.id] + if guild.id in self.birthday_tasks: + del self.birthday_tasks[guild.id] await self.config.guild(guild).scheduled_tasks.clear_raw(str(member.id)) async def reload_scheduled_tasks(self): """Reload and reschedule tasks from the configuration.""" - for guild in self.bot.guilds: - scheduled_tasks = await self.config.guild(guild).scheduled_tasks() - for member_id, task_info in scheduled_tasks.items(): - member = guild.get_member(int(member_id)) - if not member: - continue - role = guild.get_role(task_info["role_id"]) - if not role: - continue - remove_at = datetime.fromisoformat(task_info["remove_at"]).replace(tzinfo=ZoneInfo(await self.config.guild(guild).timezone())) - self.birthday_tasks[guild.id] = self.bot.loop.create_task(self.remove_birthday_role(guild, member, role, remove_at)) + try: + logger.info("Reloading scheduled birthday tasks") + for guild in self.bot.guilds: + try: + scheduled_tasks = await self.config.guild(guild).scheduled_tasks() + for member_id, task_info in scheduled_tasks.items(): + try: + member = guild.get_member(int(member_id)) + if not member: + logger.warning(f"Member {member_id} not found in guild {guild.id}, removing task") + await self.config.guild(guild).scheduled_tasks.clear_raw(member_id) + continue + + role = guild.get_role(task_info["role_id"]) + if not role: + logger.warning(f"Role {task_info['role_id']} not found in guild {guild.id}, removing task") + await self.config.guild(guild).scheduled_tasks.clear_raw(member_id) + continue + + remove_at = datetime.fromisoformat(task_info["remove_at"]).replace( + tzinfo=ZoneInfo(await self.config.guild(guild).timezone())) + + if datetime.now(remove_at.tzinfo) >= remove_at: + try: + await member.remove_roles(role, reason="Birthday role duration expired (reload)") + logger.info(f"Removed expired birthday role from {member_id} in guild {guild.id}") + except discord.Forbidden: + logger.error(f"Failed to remove birthday role from {member_id} in guild {guild.id}: Insufficient permissions") + except discord.HTTPException as e: + logger.error(f"Failed to remove birthday role from {member_id} in guild {guild.id}: {str(e)}") + finally: + await self.config.guild(guild).scheduled_tasks.clear_raw(member_id) + else: + self.birthday_tasks[guild.id] = self.bot.loop.create_task( + self.remove_birthday_role(guild, member, role, remove_at)) + logger.info(f"Rescheduled birthday role removal for {member_id} in guild {guild.id}") + except Exception as e: + logger.error(f"Error processing task for member {member_id} in guild {guild.id}: {str(e)}", exc_info=True) + except Exception as e: + logger.error(f"Error processing guild {guild.id} during reload: {str(e)}", exc_info=True) + except Exception as e: + logger.error(f"Error reloading scheduled tasks: {str(e)}", exc_info=True)