diff --git a/src/cogs/moderation/infraction.py b/src/cogs/moderation/infraction.py index d6c4f1c..c7d1096 100644 --- a/src/cogs/moderation/infraction.py +++ b/src/cogs/moderation/infraction.py @@ -30,8 +30,12 @@ def __init__(self, bot: APBot) -> None: def has_mod_role(self, member: Member) -> bool: allowed_role_names = {"Trial Chat Moderator", "Chat Moderator", "Admin"} allowed_role_ids = { - conf.get("bot_staff_role_id"), - conf.get("special_perms_role_id"), + role_id + for role_id in ( + conf.get("bot_staff_role_id"), + conf.get("special_perms_role_id"), + ) + if role_id is not None } perms = getattr(member, "guild_permissions", None) diff --git a/src/cogs/study.py b/src/cogs/study.py index 5d998bb..7eff8ce 100644 --- a/src/cogs/study.py +++ b/src/cogs/study.py @@ -3,16 +3,49 @@ import asyncio import nextcord -from nextcord import Interaction, Embed, Attachment, Message, slash_command, SlashOption -from nextcord.ext import commands -from typing import Union +from nextcord import Interaction, Embed, Attachment, slash_command, SlashOption +from nextcord.ext import commands, tasks from cogs.utils import convert_time blue = nextcord.Color.teal() -from config_handler import Config -config_path = "config.json" -conf = Config(config_path) +from app_config import get_command_guild_ids, load_optional_config +conf = load_optional_config() +COMMAND_GUILD_IDS = get_command_guild_ids(conf) + +MIN_STUDY_SECONDS = 60 * 10 +MAX_STUDY_SECONDS = 60 * 60 * 24 * 7 + + +def validate_study_duration(duration: str): + duration_seconds = convert_time(duration) + + if isinstance(duration_seconds, str): + return None, duration_seconds + + if duration_seconds < MIN_STUDY_SECONDS: + return None, "Please set a duration greater than or equal to 10 minutes." + + if duration_seconds > MAX_STUDY_SECONDS: + return None, "Please set a duration less than or equal to one week." + + return duration_seconds, None + + +def get_study_delay_seconds(study_end: int, now: int | None = None) -> int: + current_time = int(time.time()) if now is None else now + return max(0, int(study_end) - current_time) + + +async def send_study_role_dm(member, study_end: int) -> bool: + try: + await member.send( + f"You enabled Study Mode. The Study role will be removed at ." + ) + except (nextcord.Forbidden, nextcord.HTTPException): + return False + + return True class QuestionConfirm(nextcord.ui.View): @@ -51,9 +84,13 @@ def __init__(self, bot: commands.Bot): self.bot = bot self.last_topic_update = {} self.potd_cooldowns = {} + + def cog_unload(self) -> None: + if self.expire_study_roles.is_running(): + self.expire_study_roles.cancel() - @slash_command(name="question", description="Ask a question in a subject channel.", guild_ids=[conf.get("guild_id")]) + @slash_command(name="question", description="Ask a question in a subject channel.", guild_ids=COMMAND_GUILD_IDS) async def question( self, inter: Interaction, @@ -106,7 +143,7 @@ async def question( - @slash_command(name="potd", description="Make a problem-of-the-day for a subject channel.", guild_ids=[conf.get("guild_id")]) + @slash_command(name="potd", description="Make a problem-of-the-day for a subject channel.", guild_ids=COMMAND_GUILD_IDS) async def potd( self, inter: Interaction, @@ -174,21 +211,57 @@ async def potd( + @tasks.loop(minutes=1) + async def expire_study_roles(self) -> None: + try: + study_users = await self.bot.db.study.get_all() + except Exception as exc: + print(f"[Study] Failed to fetch study role expirations: {exc}") + return + + now = int(time.time()) + for user_id, study_end in study_users.items(): + try: + if int(study_end) <= now: + await self.remove_study_role(int(user_id)) + except Exception: + continue + + @expire_study_roles.before_loop + async def before_expire_study_roles(self) -> None: + await self.bot.wait_until_ready() + + @commands.Cog.listener() + async def on_ready(self) -> None: + if not self.expire_study_roles.is_running(): + self.expire_study_roles.start() + async def remove_study_role(self, user_id: int) -> None: + guild_id = self.bot.config.get("guild_id") + try: - guild = await self.bot.fetch_guild(self.bot.config.get("guild_id")) + guild = await self.bot.getch_guild(guild_id) user = await self.bot.getch_member(guild.id, user_id) study_role = nextcord.utils.get(guild.roles, name="Study") - except: + except Exception: + return + + if not study_role: return - if not user or not study_role: + if not user: + await self.bot.db.study.delete_user(user_id) return + if study_role.id in [role.id for role in user.roles]: + try: + await user.remove_roles(study_role, reason="Study role duration expired.") + except (nextcord.Forbidden, nextcord.HTTPException): + return + await self.bot.db.study.delete_user(user_id) - await user.remove_roles(study_role) - @slash_command(name="study", description="Prevent yourself from viewing unhelpful channels.", guild_ids=[conf.get("guild_id")]) + @slash_command(name="study", description="Prevent yourself from viewing unhelpful channels.", guild_ids=COMMAND_GUILD_IDS) async def study( self, inter: Interaction, @@ -197,32 +270,36 @@ async def study( await inter.response.defer(ephemeral=True) resp = await inter.send("Validating time input...", ephemeral=True) - duration: Union[str, int] = convert_time(duration) - if isinstance(duration, str): - return await resp.edit(content=duration) + duration_seconds, error = validate_study_duration(duration) + if error: + return await resp.edit(content=error) - if duration < 60 * 10: - return await resp.edit("Please set a duration greater than 10 minutes!") - if duration > 60 * 60 * 24 * 7: - return await resp.edit("Please set a duration lesser than a week") + await resp.edit(content="Performing actions...") + study_end = int(time.time()) + duration_seconds - await resp.edit("Performing actions...") - study_end = int(time.time()) + duration + guild = await self.bot.getch_guild(self.bot.config.get("guild_id")) + if not guild: + return await resp.edit(content="Server not found.") - guild = await self.bot.fetch_guild(self.bot.config.get("guild_id")) study_role = nextcord.utils.get(guild.roles, name="Study") - if study_role.id in [role.id for role in inter.user.roles]: - await resp.edit("Removing role...") - await self.remove_study_role(inter.user.id) + if study_role is None: + return await resp.edit(content="Study role not found.") - await inter.user.add_roles(study_role) - await resp.edit("Updating database...") + user_role_ids = {role.id for role in inter.user.roles} + if study_role.id not in user_role_ids: + await inter.user.add_roles(study_role, reason="Study role enabled by user.") + await resp.edit(content="Updating database...") await self.bot.db.study.set_time(inter.user.id, study_end) - self.bot.loop.call_later(duration, asyncio.create_task, self.remove_study_role(inter.user.id)) + user_id = inter.user.id + self.bot.loop.call_later( + get_study_delay_seconds(study_end), + lambda user_id=user_id: asyncio.create_task(self.remove_study_role(user_id)), + ) + await send_study_role_dm(inter.user, study_end) await resp.edit( - f"The study role will be removed at ." + content=f"The study role will be removed at ." ) class ConfirmDeny(nextcord.ui.View): diff --git a/tests/test_study_helpers.py b/tests/test_study_helpers.py new file mode 100644 index 0000000..b13767f --- /dev/null +++ b/tests/test_study_helpers.py @@ -0,0 +1,38 @@ +import asyncio + +from cogs.study import MAX_STUDY_SECONDS, MIN_STUDY_SECONDS, get_study_delay_seconds, send_study_role_dm, validate_study_duration + + +class FakeMember: + def __init__(self): + self.sent = None + + async def send(self, content): + self.sent = content + + +def test_validate_study_duration_accepts_bounds(): + assert validate_study_duration("10m") == (MIN_STUDY_SECONDS, None) + assert validate_study_duration("168h") == (MAX_STUDY_SECONDS, None) + assert validate_study_duration("1w") == (MAX_STUDY_SECONDS, None) + + +def test_validate_study_duration_rejects_out_of_range(): + assert validate_study_duration("9m")[0] is None + assert validate_study_duration("169h")[0] is None + + +def test_send_study_role_dm_sends_message(): + member = FakeMember() + + result = asyncio.run(send_study_role_dm(member, 1234567890)) + + assert result is True + assert "Study Mode" in member.sent + assert "" in member.sent + + +def test_get_study_delay_seconds(): + assert get_study_delay_seconds(1000, now=900) == 100 + assert get_study_delay_seconds(1000, now=1000) == 0 + assert get_study_delay_seconds(1000, now=1100) == 0