From fdf12c6d2b2f3ab5ae335e2913a714cbeac2ff30 Mon Sep 17 00:00:00 2001 From: mbaruh Date: Sat, 13 Feb 2021 16:14:56 +0200 Subject: Add option to schedule threshold reset Added optional argument to defcon threshold to specify for how long it should be on. The notifier will now run only when there is no expiry date specified. --- bot/exts/moderation/defcon.py | 62 ++++++++++++++++++++++++++++++++----------- 1 file changed, 46 insertions(+), 16 deletions(-) diff --git a/bot/exts/moderation/defcon.py b/bot/exts/moderation/defcon.py index 82aaf5714..8c21a7327 100644 --- a/bot/exts/moderation/defcon.py +++ b/bot/exts/moderation/defcon.py @@ -4,19 +4,20 @@ import logging from collections import namedtuple from datetime import datetime from enum import Enum -from typing import Union +from typing import Optional, Union from async_rediscache import RedisCache from dateutil.relativedelta import relativedelta -from discord import Colour, Embed, Member +from discord import Colour, Embed, Member, User from discord.ext import tasks from discord.ext.commands import Cog, Context, group, has_any_role from bot.bot import Bot from bot.constants import Channels, Colours, Emojis, Event, Icons, MODERATION_ROLES, Roles -from bot.converters import DurationDelta +from bot.converters import DurationDelta, Expiry from bot.exts.moderation.modlog import ModLog from bot.utils.messages import format_user +from bot.utils.scheduling import Scheduler from bot.utils.time import humanize_delta, parse_duration_string log = logging.getLogger(__name__) @@ -60,6 +61,8 @@ class Defcon(Cog): self.threshold = relativedelta(days=0) self.expiry = None + self.scheduler = Scheduler(self.__class__.__name__) + self.bot.loop.create_task(self._sync_settings()) @property @@ -79,11 +82,15 @@ class Defcon(Cog): try: settings = await self.redis_cache.to_dict() self.threshold = parse_duration_string(settings["threshold"]) + self.expiry = datetime.fromisoformat(settings["expiry"]) if settings["expiry"] else None except Exception: log.exception("Unable to get DEFCON settings!") await self.channel.send(f"<@&{Roles.moderators}> **WARNING**: Unable to get DEFCON settings!") else: + if self.expiry: + self.scheduler.schedule_at(self.expiry, 0, self._remove_threshold()) + self._update_notifier() log.info(f"DEFCON synchronized: {humanize_delta(self.threshold)}") @@ -95,7 +102,7 @@ class Defcon(Cog): if self.threshold > relativedelta(days=0): now = datetime.utcnow() - if now - member.created_at < self.threshold: + if now - member.created_at < self.threshold: # TODO log.info(f"Rejecting user {member}: Account is too new") message_sent = False @@ -104,7 +111,7 @@ class Defcon(Cog): await member.send(REJECTION_MESSAGE.format(user=member.mention)) message_sent = True - except Exception: # TODO + except Exception: log.exception(f"Unable to send rejection message to user: {member}") await member.kick(reason="DEFCON active, user is too new") @@ -132,17 +139,22 @@ class Defcon(Cog): """Check the current status of DEFCON mode.""" embed = Embed( colour=Colour.blurple(), title="DEFCON Status", - description=f"**Threshold:** {humanize_delta(self.threshold)}" + description=f""" + **Threshold:** {humanize_delta(self.threshold)} + **Expires in:** {humanize_delta(relativedelta(self.expiry, datetime.utcnow())) if self.expiry else "-"} + """ ) await ctx.send(embed=embed) @defcon_group.command(aliases=('t',)) - async def threshold(self, ctx: Context, threshold: Union[DurationDelta, int]) -> None: + async def threshold( + self, ctx: Context, threshold: Union[DurationDelta, int], expiry: Optional[Expiry] = None + ) -> None: """Set how old an account must be to join the server.""" if isinstance(threshold, int): threshold = relativedelta(days=threshold) - await self._defcon_action(ctx, threshold=threshold) + await self._defcon_action(ctx.author, threshold=threshold, expiry=expiry) @defcon_group.command() async def shutdown(self, ctx: Context) -> None: @@ -172,28 +184,45 @@ class Defcon(Cog): await self.channel.edit(topic=new_topic) @redis_cache.atomic_transaction - async def _defcon_action(self, ctx: Context, threshold: relativedelta) -> None: + async def _defcon_action(self, author: User, threshold: relativedelta, expiry: Optional[Expiry] = None) -> None: """Providing a structured way to do a defcon action.""" self.threshold = threshold + if threshold == relativedelta(days=0): # If the threshold is 0, we don't need to schedule anything + expiry = None + self.expiry = expiry + + # Either way, we cancel the old task. + self.scheduler.cancel_all() + if self.expiry is not None: + self.scheduler.schedule_at(expiry, 0, self._remove_threshold()) await self.redis_cache.update( { 'threshold': Defcon._stringify_relativedelta(self.threshold), + 'expiry': expiry.isoformat() if expiry else 0 } ) self._update_notifier() action = Action.DURATION_UPDATE - await ctx.send( + expiry_message = "" + if expiry: + expiry_message = f"for the next {humanize_delta(relativedelta(expiry, datetime.utcnow()))}" + + await self.channel.send( f"{action.value.emoji} DEFCON threshold updated; accounts must be " - f"{humanize_delta(self.threshold)} old to join the server." + f"{humanize_delta(self.threshold)} old to join the server {expiry_message}." ) - await self._send_defcon_log(action, ctx.author) + await self._send_defcon_log(action, author) await self._update_channel_topic() self._log_threshold_stat(threshold) + async def _remove_threshold(self) -> None: + """Resets the threshold back to 0.""" + await self._defcon_action(self.bot.user, relativedelta(days=0)) + @staticmethod def _stringify_relativedelta(delta: relativedelta) -> str: """Convert a relativedelta object to a duration string.""" @@ -206,7 +235,7 @@ class Defcon(Cog): threshold_days = (utcnow + threshold - utcnow).total_seconds() / SECONDS_IN_DAY self.bot.stats.gauge("defcon.threshold", threshold_days) - async def _send_defcon_log(self, action: Action, actor: Member) -> None: + async def _send_defcon_log(self, action: Action, actor: User) -> None: """Send log message for DEFCON action.""" info = action.value log_msg: str = ( @@ -219,11 +248,11 @@ class Defcon(Cog): def _update_notifier(self) -> None: """Start or stop the notifier according to the DEFCON status.""" - if self.threshold != relativedelta(days=0) and not self.defcon_notifier.is_running(): + if self.threshold != relativedelta(days=0) and self.expiry is None and not self.defcon_notifier.is_running(): log.info("DEFCON notifier started.") self.defcon_notifier.start() - elif self.threshold == relativedelta(days=0) and self.defcon_notifier.is_running(): + elif (self.threshold == relativedelta(days=0) or self.expiry is not None) and self.defcon_notifier.is_running(): log.info("DEFCON notifier stopped.") self.defcon_notifier.cancel() @@ -237,9 +266,10 @@ class Defcon(Cog): return (await has_any_role(*MODERATION_ROLES).predicate(ctx)) and ctx.channel == self.channel def cog_unload(self) -> None: - """Cancel the notifer task when the cog unloads.""" + """Cancel the notifer and threshold removal tasks when the cog unloads.""" log.trace("Cog unload: canceling defcon notifier task.") self.defcon_notifier.cancel() + self.scheduler.cancel_all() def setup(bot: Bot) -> None: -- cgit v1.2.3