diff options
Diffstat (limited to '')
| -rw-r--r-- | bot/exts/moderation/modpings.py | 96 | 
1 files changed, 94 insertions, 2 deletions
diff --git a/bot/exts/moderation/modpings.py b/bot/exts/moderation/modpings.py index 1ad5005de..bafd40580 100644 --- a/bot/exts/moderation/modpings.py +++ b/bot/exts/moderation/modpings.py @@ -1,8 +1,9 @@ +import asyncio  import datetime  import logging  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 @@ -13,6 +14,8 @@ from bot.utils.scheduling import Scheduler  log = logging.getLogger(__name__) +MINIMUM_WORK_LIMIT = 16 +  class ModPings(Cog):      """Commands for a moderator to turn moderator pings on and off.""" @@ -22,14 +25,21 @@ 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.reschedule_task = self.bot.loop.create_task(self.reschedule_roles(), name="mod-pings-reschedule") +        self.modpings_schedule_task = self.bot.loop.create_task(self.reschedule_modpings_schedule())      async def reschedule_roles(self) -> None:          """Reschedule moderators role re-apply times.""" @@ -55,6 +65,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: +            start_timestamp, work_time = schedule.split("|") +            start = datetime.datetime.fromtimestamp(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(minutes=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}.") @@ -126,12 +183,47 @@ class ModPings(Cog):          await ctx.send(f"{Emojis.check_mark} Moderators role has been re-applied.") +    @modpings_group.command(name='schedule') +    @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=MINIMUM_WORK_LIMIT): +            await ctx.send( +                f":x: {ctx.author.mention} You need to have the role on for " +                f"a minimum of {MINIMUM_WORK_LIMIT} hours!" +            ) +            return + +        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}") + +        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: %I:%M%p} to {end: %I:%M%p} UTC Timing!" +        ) +      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."""  |