diff options
| -rw-r--r-- | bot/exts/moderation/modpings.py | 118 |
1 files changed, 116 insertions, 2 deletions
diff --git a/bot/exts/moderation/modpings.py b/bot/exts/moderation/modpings.py index a7ccb8162..6cc46ad26 100644 --- a/bot/exts/moderation/modpings.py +++ b/bot/exts/moderation/modpings.py @@ -1,7 +1,8 @@ +import asyncio import datetime from async_rediscache import RedisCache -from dateutil.parser import isoparse +from dateutil.parser import isoparse, parse from discord import Embed, Member from discord.ext.commands import Cog, Context, group, has_any_role @@ -14,6 +15,8 @@ from bot.utils.scheduling import Scheduler log = get_logger(__name__) +MAXIMUM_WORK_LIMIT = 16 + class ModPings(Cog): """Commands for a moderator to turn moderator pings on and off.""" @@ -23,13 +26,23 @@ class ModPings(Cog): # The cache's values are the times when the role should be re-applied to them, stored in ISO format. pings_off_mods = RedisCache() + # RedisCache[discord.Member.id, 'start timestamp|total worktime in seconds'] + # The cache's keys are mod's ID + # The cache's values are their pings on schedule timestamp and the total seconds (work time) until pings off + modpings_schedule = RedisCache() + def __init__(self, bot: Bot): self.bot = bot - self._role_scheduler = Scheduler(self.__class__.__name__) + self._role_scheduler = Scheduler("ModPingsOnOff") + self._modpings_scheduler = Scheduler("ModPingsSchedule") self.guild = None self.moderators_role = None + self.modpings_schedule_task = scheduling.create_task( + self.reschedule_modpings_schedule(), + event_loop=self.bot.loop + ) self.reschedule_task = scheduling.create_task( self.reschedule_roles(), name="mod-pings-reschedule", @@ -60,6 +73,53 @@ class ModPings(Cog): expiry = isoparse(pings_off[mod.id]).replace(tzinfo=None) self._role_scheduler.schedule_at(expiry, mod.id, self.reapply_role(mod)) + async def reschedule_modpings_schedule(self) -> None: + """Reschedule moderators schedule ping.""" + await self.bot.wait_until_guild_available() + schedule_cache = await self.modpings_schedule.to_dict() + + log.info("Scheduling modpings schedule for applicable moderators found in cache.") + for mod_id, schedule in schedule_cache.items(): + start_timestamp, work_time = schedule.split("|") + start = datetime.datetime.fromtimestamp(float(start_timestamp)) + + mod = self.bot.fetch_user(mod_id) + self._modpings_scheduler.schedule_at( + start, + mod_id, + self.add_role_schedule(mod, work_time, start) + ) + + async def remove_role_schedule(self, mod: Member, work_time: int, schedule_start: datetime.datetime) -> None: + """Removes the moderator's role to the given moderator.""" + log.trace(f"Removing moderator role to mod with ID {mod.id}") + await mod.remove_roles(self.moderators_role, reason="Moderator schedule time expired.") + + # Remove the task before scheduling it again + self._modpings_scheduler.cancel(mod.id) + + # Add the task again + log.trace(f"Adding mod pings schedule task again for mod with ID {mod.id}") + schedule_start += datetime.timedelta(days=1) + self._modpings_scheduler.schedule_at( + schedule_start, + mod.id, + self.add_role_schedule(mod, work_time, schedule_start) + ) + + async def add_role_schedule(self, mod: Member, work_time: int, schedule_start: datetime.datetime) -> None: + """Adds the moderator's role to the given moderator.""" + # If the moderator has pings off, then skip adding role + if mod.id in await self.pings_off_mods.to_dict(): + log.trace(f"Skipping adding moderator role to mod with ID {mod.id} - found in pings off cache.") + else: + log.trace(f"Applying moderator role to mod with ID {mod.id}") + await mod.add_roles(self.moderators_role, reason="Moderator schedule time started!") + + log.trace(f"Sleeping for {work_time} seconds, worktime for mod with ID {mod.id}") + await asyncio.sleep(work_time) + await self.remove_role_schedule(mod, work_time, schedule_start) + async def reapply_role(self, mod: Member) -> None: """Reapply the moderator's role to the given moderator.""" log.trace(f"Re-applying role to mod with ID {mod.id}.") @@ -131,12 +191,66 @@ class ModPings(Cog): await ctx.send(f"{Emojis.check_mark} Moderators role has been re-applied.") + @modpings_group.group( + name='schedule', + aliases=('s',), + invoke_without_command=True + ) + @has_any_role(*MODERATION_ROLES) + async def schedule_modpings(self, ctx: Context, start: str, end: str) -> None: + """Schedule modpings role to be added at <start> and removed at <end> everyday at UTC time!""" + start, end = parse(start), parse(end) + + if end < start: + end += datetime.timedelta(days=1) + + if (end - start) > datetime.timedelta(hours=MAXIMUM_WORK_LIMIT): + await ctx.send( + f":x: {ctx.author.mention} You can't have the modpings role for" + f" more than {MAXIMUM_WORK_LIMIT} hours!" + ) + return + + if start < datetime.datetime.utcnow(): + # The datetime has already gone for the day, so make it tomorrow + # otherwise the scheduler would schedule it immediately + start += datetime.timedelta(days=1) + + start, end = start.replace(tzinfo=None), end.replace(tzinfo=None) + work_time = (end - start).total_seconds() + + await self.modpings_schedule.set(ctx.author.id, f"{start.timestamp()}|{work_time}") + + if ctx.author.id in self._modpings_scheduler: + self._modpings_scheduler.cancel(ctx.author.id) + + self._modpings_scheduler.schedule_at( + start, + ctx.author.id, + self.add_role_schedule(ctx.author, work_time, start) + ) + + await ctx.send( + f"{Emojis.ok_hand} {ctx.author.mention} Scheduled mod pings from " + f"{start: %H:%M} to {end: %H:%M} UTC Timing!" + ) + + @schedule_modpings.command(name='delete', aliases=('del', 'd')) + async def modpings_schedule_delete(self, ctx: Context) -> None: + """Delete your modpings schedule.""" + self._modpings_scheduler.cancel(ctx.author.id) + await self.modpings_schedule.delete(ctx.author.id) + await ctx.send(f"{Emojis.ok_hand} {ctx.author.mention} Deleted your modpings schedule!") + def cog_unload(self) -> None: """Cancel role tasks when the cog unloads.""" log.trace("Cog unload: canceling role tasks.") self.reschedule_task.cancel() self._role_scheduler.cancel_all() + self.modpings_schedule_task.cancel() + self._modpings_scheduler.cancel_all() + def setup(bot: Bot) -> None: """Load the ModPings cog.""" |