From 78e0f5bc91ff10e69014bb263f5408f0358836f1 Mon Sep 17 00:00:00 2001 From: Ryuga Date: Fri, 1 May 2026 00:12:55 +0530 Subject: [PATCH] Schedule mod mail pings --- bot/bot.py | 5 +- bot/cogs/tortoise_dm.py | 114 +++++++++++++++++++++++++++++++++++++++- bot/manager.py | 41 +++++++++++++++ 3 files changed, 158 insertions(+), 2 deletions(-) diff --git a/bot/bot.py b/bot/bot.py index ef44c90..d214605 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -15,7 +15,7 @@ from bot.api_client import TortoiseAPI from bot.constants import error_log_channel_id, system_log_channel_id, github_repo_link from bot.manager import ( - Database, ProgressionManager, AFKManager, PointsManager, RetentionManager, TeamManager, GiveawayManager + Database, ProgressionManager, AFKManager, PointsManager, RetentionManager, TeamManager, GiveawayManager, DutyManager ) from bot.utils.embed_handler import simple_embed from bot.utils.error_handler import TortoiseCommandTree @@ -75,6 +75,7 @@ def __init__(self, prefix="t.", *args, **kwargs): self.retention_manager = None self.team_manager = None self.giveaway_manager = None + self.duty_manager = None self._sys_log_channel = None @property @@ -140,6 +141,7 @@ async def setup_hook(self): self.retention_manager = RetentionManager(self.db) self.team_manager = TeamManager(self.db) self.giveaway_manager = GiveawayManager(self.db) + self.duty_manager = DutyManager(self.db) await self.progression_manager.setup() await self.afk_manager.setup() @@ -147,6 +149,7 @@ async def setup_hook(self): await self.retention_manager.setup() await self.team_manager.setup() await self.giveaway_manager.setup() + await self.duty_manager.setup() await self.load_extensions() # await self.reload_tortoise_meta_cache() diff --git a/bot/cogs/tortoise_dm.py b/bot/cogs/tortoise_dm.py index fb2e2f2..6870b58 100644 --- a/bot/cogs/tortoise_dm.py +++ b/bot/cogs/tortoise_dm.py @@ -1,12 +1,16 @@ +import datetime import logging from io import StringIO from typing import Union from asyncio import TimeoutError +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import discord -from discord.ext import commands +from discord.ext import commands, tasks +from discord import app_commands from bot import constants +from bot.utils.checks import check_if_tortoise_staff from bot.utils.cooldown import CoolDown from bot.utils.message_logger import MessageLogger from bot.utils.embed_handler import authored, failure, success, info, create_suggestion_msg, authored_sm @@ -22,6 +26,57 @@ class UnsupportedFileExtension(Exception): class UnsupportedFileEncoding(ValueError): pass +class DutyScheduleModal(discord.ui.Modal, title="Set Daily Mod Mail Schedule"): + start_time = discord.ui.TextInput( + label="Start Time (24h format)", + placeholder="e.g. 09:00", + min_length=5, max_length=5 + ) + end_time = discord.ui.TextInput( + label="End Time (24h format)", + placeholder="e.g. 17:00", + min_length=5, max_length=5 + ) + timezone = discord.ui.TextInput( + label="Your Timezone (IANA Name)", + placeholder="e.g. Europe/London, America/New_York, Asia/Kolkata", + min_length=4 + ) + + def __init__(self, cog: "TortoiseDM"): + super().__init__() + self.cog = cog + + async def on_submit(self, interaction: discord.Interaction): + tz_str = self.timezone.value.strip().replace(" ", "_") + try: + datetime.datetime.strptime(self.start_time.value, "%H:%M") + datetime.datetime.strptime(self.end_time.value, "%H:%M") + ZoneInfo(tz_str) + except ValueError: + await interaction.response.send_message(embed=failure("Invalid time format. Use HH:MM."), ephemeral=True) + return + except ZoneInfoNotFoundError: + await interaction.response.send_message(embed=failure( + "Invalid timezone. Example: 'UTC' or 'Europe/London'." + ), ephemeral=True) + return + + await self.cog.duty_manager.set_schedule( + interaction.guild.id, + interaction.user.id, + self.start_time.value, + self.end_time.value, + tz_str + ) + + await interaction.response.send_message( + embed=success( + f"Schedule set! You'll receive the pings daily between " + f"{self.start_time.value} and {self.end_time.value} ({tz_str})." + ), ephemeral=True + ) + class DMInitView(discord.ui.View): def __init__(self, cog: "TortoiseDM", user: discord.User): super().__init__(timeout=300) @@ -211,6 +266,8 @@ def mod_mail_check(msg): class TortoiseDM(commands.Cog): + mod_mail_group = app_commands.Group(name="mod_mail", description="Manage your mod mail ping schedule") + def __init__(self, bot): self.bot = bot self._tortoise_guild = None @@ -219,6 +276,7 @@ def __init__(self, bot): self._mod_mail_ping_role = None self.cool_down = CoolDown(seconds=120) self.bot.loop.create_task(self.cool_down.start()) + self.duty_manager = bot.duty_manager # Key is user id value is mod/admin id self.active_mod_mails = {} @@ -280,6 +338,9 @@ async def on_ready(self): self.staff_channel = self.bot.get_channel(constants.staff_channel_id) self.staff_applications_channel = self.bot.get_channel(constants.system_log_channel_id) + if not self.duty_automation_loop.is_running(): + self.duty_automation_loop.start() + @property def tortoise_guild(self): if self._tortoise_guild is None: @@ -304,6 +365,44 @@ def mod_mail_ping_role(self): self._mod_mail_ping_role = self.tortoise_guild.get_role(constants.mod_mail_ping_role_id) return self._mod_mail_ping_role + @tasks.loop(hours=1) + async def duty_automation_loop(self): + """Background task running every minute to check schedules.""" + await self.bot.wait_until_ready() + try: + schedules = await self.duty_manager.get_all_schedules() + now_utc = datetime.datetime.now(datetime.timezone.utc) + + for record in schedules: + guild = self.bot.get_guild(record["guild_id"]) + if not guild: continue + + member = guild.get_member(record["user_id"]) + if not member: continue + + user_tz = ZoneInfo(record["timezone"]) + local_now = now_utc.astimezone(user_tz) + current_local_hm = local_now.strftime("%H:%M") + + start = record["start_time"] + end = record["end_time"] + + if start < end: + is_duty = start <= current_local_hm < end + else: + is_duty = current_local_hm >= start or current_local_hm < end + + role = self.mod_mail_ping_role + has_role = role in member.roles + + if is_duty and not has_role: + await member.add_roles(role, reason="Scheduled duty started.") + elif not is_duty and has_role: + await member.remove_roles(role, reason="Scheduled duty ended.") + + except Exception as e: + logger.error(f"Error in duty loop: {e}") + @commands.Cog.listener() async def on_message(self, message): if message.author == self.bot.user: @@ -631,6 +730,19 @@ def _get_attachments_as_urls(cls, message: discord.Message) -> str: urls = '\n'.join(attachment.url for attachment in message.attachments) return f"\nAttachments:\n{urls}" + @mod_mail_group.command(name="ping_schedule", description="Set your daily recurring mod mail ping hours.") + @app_commands.check(check_if_tortoise_staff) + async def set_ping_schedule(self, interaction: discord.Interaction): + await interaction.response.send_modal(DutyScheduleModal(self)) + + @mod_mail_group.command(name="stop_ping", description="Remove your automatic ping schedule.") + @app_commands.check(check_if_tortoise_staff) + async def stop_duty_schedule(self, interaction: discord.Interaction): + await self.duty_manager.remove_schedule(interaction.guild.id, interaction.user.id) + + if self.mod_mail_ping_role in interaction.user.roles: + await interaction.user.remove_roles(self.mod_mail_ping_role, reason="Schedule deleted.") + await interaction.response.send_message(embed=success("Your ping schedule has been deleted."), ephemeral=True) async def setup(bot): await bot.add_cog(TortoiseDM(bot)) diff --git a/bot/manager.py b/bot/manager.py index dc68f6b..3b9c83d 100644 --- a/bot/manager.py +++ b/bot/manager.py @@ -763,3 +763,44 @@ async def delete_giveaway(self, message_id: int): message_id, ) +class DutyManager: + def __init__(self, db: Database): + self.db = db + + async def setup(self): + await self.db.pool.execute( + """ + CREATE TABLE IF NOT EXISTS duty_schedules ( + guild_id BIGINT NOT NULL, + user_id BIGINT NOT NULL, + start_time TEXT NOT NULL, + end_time TEXT NOT NULL, + timezone TEXT NOT NULL, + PRIMARY KEY (guild_id, user_id) + ) + """ + ) + + async def set_schedule(self, guild_id: int, user_id: int, start: str, end: str, tz: str): + await self.db.pool.execute( + """ + INSERT INTO duty_schedules (guild_id, user_id, start_time, end_time, timezone) + VALUES ($1, $2, $3, $4, $5) + ON CONFLICT (guild_id, user_id) + DO UPDATE SET + start_time = EXCLUDED.start_time, + end_time = EXCLUDED.end_time, + timezone = EXCLUDED.timezone + """, + guild_id, user_id, start, end, tz + ) + + async def remove_schedule(self, guild_id: int, user_id: int): + await self.db.pool.execute( + "DELETE FROM duty_schedules WHERE guild_id = $1 AND user_id = $2", + guild_id, user_id + ) + + async def get_all_schedules(self): + return await self.db.pool.fetch("SELECT * FROM duty_schedules") +