From e0c3b0d0e1d9986cd31f0271caa3032149bb69cf Mon Sep 17 00:00:00 2001 From: petabyte-imo Date: Wed, 4 Sep 2024 13:55:28 +0100 Subject: [PATCH 01/11] Added a tempban command and a function that gets all guilds from the db --- tux/cogs/moderation/tempban.py | 120 ++++++++++++++++++++++++++++++ tux/database/controllers/guild.py | 3 + 2 files changed, 123 insertions(+) create mode 100644 tux/cogs/moderation/tempban.py diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py new file mode 100644 index 000000000..828c09e90 --- /dev/null +++ b/tux/cogs/moderation/tempban.py @@ -0,0 +1,120 @@ +from datetime import UTC, datetime + +import discord +from discord.ext import commands, tasks +from loguru import logger + +from prisma.enums import CaseType +from tux.bot import Tux +from tux.utils import checks +from tux.utils.flags import TempBanFlags, generate_usage +from tux.utils.functions import parse_time_string + +from . import ModerationCogBase + + +class TempBan(ModerationCogBase): + def __init__(self, bot: Tux) -> None: + super().__init__(bot) + self.tempban.usage = generate_usage(self.tempban, TempBanFlags) + self.tempban_check.start() + + @commands.hybrid_command(name="tempban", aliases=["tb"]) + @commands.guild_only() + @checks.has_pl(3) + async def tempban( + self, + ctx: commands.Context[Tux], + member: discord.Member, + *, + flags: TempBanFlags, + ) -> None: + """ + Ban a member from the server. + + Parameters + ---------- + ctx : commands.Context[Tux] + The context in which the command is being invoked. + member : discord.Member + The member to ban. + flags : TempBanFlags + The flags for the command. (reason: str, purge_days: int, silent: bool, expires_at: int The number of days the ban will last for.) + + Raises + ------ + discord.Forbidden + If the bot is unable to ban the user. + discord.HTTPException + If an error occurs while banning the user. + """ + + if ctx.guild is None: + logger.warning("Ban command used outside of a guild context.") + return + + moderator = ctx.author + duration = parse_time_string(f"{flags.expires_at}d") + expires_at = datetime.now(UTC) + duration + # Debug Print + logger.debug(f"Temp Ban Duration: {duration}\nTemp Ban Time Expires: {expires_at}") + if not await self.check_conditions(ctx, member, moderator, "temporarily ban"): + return + + try: + await self.send_dm(ctx, flags.silent, member, flags.reason, action="temporarily banned") + await ctx.guild.ban(member, reason=flags.reason, delete_message_days=flags.purge_days) + + except (discord.Forbidden, discord.HTTPException) as e: + logger.error(f"Failed to temporarily ban {member}. {e}") + await ctx.send(f"Failed to temporarily ban {member}. {e}", delete_after=30, ephemeral=True) + return + + case = await self.db.case.insert_case( + case_user_id=member.id, + case_moderator_id=ctx.author.id, + case_type=CaseType.TEMPBAN, + case_reason=flags.reason, + guild_id=ctx.guild.id, + case_expires_at=expires_at, + ) + + await self.handle_case_response(ctx, CaseType.TEMPBAN, case.case_number, flags.reason, member) + + @tasks.loop(minutes=30) + async def tempban_check(self) -> None: + # Fetch all guilds and fetch all tempbans for each guild's ID + guilds = await self.db.guild.get_all_guilds() + tempbanss = [await self.db.case.get_all_cases_by_type(guild.guild_id, CaseType.TEMPBAN) for guild in guilds] + + for tempbans in tempbanss: + for tempban in tempbans: + assert tempban.case_expires_at is not None + if tempban.case_expires_at < datetime.now(UTC): + guild = self.bot.get_guild(tempban.guild_id) + if guild is None: + try: + guild = await self.bot.fetch_guild(tempban.guild_id) + + except (discord.Forbidden, discord.HTTPException) as e: + logger.error( + f"Failed to unban user with ID {tempban.case_user_id} | Case number {tempban.case_number} | Issue: Failed to get guild with ID {tempban.guild_id}. {e}", + ) + return + else: + try: + assert isinstance(tempban.case_user_id, int) + guild_bans = guild.bans() + async for ban_entry in guild_bans: + if ban_entry.user.id == tempban.case_user_id: + await guild.unban(ban_entry.user) + + except (discord.Forbidden, discord.HTTPException) as e: + logger.error( + f"Failed to unban user with ID {tempban.case_user_id} | Case number {tempban.case_number} Issue: Failed to unban user. {e}", + ) + return + + +async def setup(bot: Tux) -> None: + await bot.add_cog(TempBan(bot)) diff --git a/tux/database/controllers/guild.py b/tux/database/controllers/guild.py index 411a4b5c6..242570eeb 100644 --- a/tux/database/controllers/guild.py +++ b/tux/database/controllers/guild.py @@ -14,3 +14,6 @@ async def insert_guild_by_id(self, guild_id: int) -> Guild: async def delete_guild_by_id(self, guild_id: int) -> None: await self.table.delete(where={"guild_id": guild_id}) + + async def get_all_guilds(self) -> list[Guild]: + return await self.table.find_many() \ No newline at end of file From f183569bdd0f8cecf34d184c94d4330376fb5332 Mon Sep 17 00:00:00 2001 From: petabyte-imo Date: Wed, 4 Sep 2024 14:27:59 +0100 Subject: [PATCH 02/11] Fixed the tempban_check looped task and added some comments --- tux/cogs/moderation/tempban.py | 71 +++++++++++++++++++--------------- 1 file changed, 39 insertions(+), 32 deletions(-) diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index 828c09e90..5940b793e 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -54,10 +54,9 @@ async def tempban( return moderator = ctx.author - duration = parse_time_string(f"{flags.expires_at}d") + duration = parse_time_string(f"{flags.expires_at}m") expires_at = datetime.now(UTC) + duration - # Debug Print - logger.debug(f"Temp Ban Duration: {duration}\nTemp Ban Time Expires: {expires_at}") + if not await self.check_conditions(ctx, member, moderator, "temporarily ban"): return @@ -85,35 +84,43 @@ async def tempban( async def tempban_check(self) -> None: # Fetch all guilds and fetch all tempbans for each guild's ID guilds = await self.db.guild.get_all_guilds() - tempbanss = [await self.db.case.get_all_cases_by_type(guild.guild_id, CaseType.TEMPBAN) for guild in guilds] - - for tempbans in tempbanss: - for tempban in tempbans: - assert tempban.case_expires_at is not None - if tempban.case_expires_at < datetime.now(UTC): - guild = self.bot.get_guild(tempban.guild_id) - if guild is None: - try: - guild = await self.bot.fetch_guild(tempban.guild_id) - - except (discord.Forbidden, discord.HTTPException) as e: - logger.error( - f"Failed to unban user with ID {tempban.case_user_id} | Case number {tempban.case_number} | Issue: Failed to get guild with ID {tempban.guild_id}. {e}", - ) - return - else: - try: - assert isinstance(tempban.case_user_id, int) - guild_bans = guild.bans() - async for ban_entry in guild_bans: - if ban_entry.user.id == tempban.case_user_id: - await guild.unban(ban_entry.user) - - except (discord.Forbidden, discord.HTTPException) as e: - logger.error( - f"Failed to unban user with ID {tempban.case_user_id} | Case number {tempban.case_number} Issue: Failed to unban user. {e}", - ) - return + tempbans = [await self.db.case.get_all_cases_by_type(guild.guild_id, CaseType.TEMPBAN) for guild in guilds] + # Here, we have 3 nested for loops because for some odd reason, tempbans is a list of a list of lists, very confusing ikr + + + for tempban in tempbans: + for cases in tempbans: + for case in cases: + + # Check if the case has expired + if case.case_expires_at < datetime.now(UTC): + # Get the guild, if that doesnt work then fetch it instead + + guild = self.bot.get_guild(case.guild_id) + if guild is None: + + try: + guild = await self.bot.fetch_guild(case.guild_id) + + except (discord.Forbidden, discord.HTTPException) as e: + logger.error( + f"Failed to unban user with ID {case.case_user_id} | Case number {case.case_number} | Issue: Failed to get guild with ID {case.guild_id}. {e}", + ) + return + else: + try: + # Unban the user + + guild_bans = guild.bans() + async for ban_entry in guild_bans: + if ban_entry.user.id == case.case_user_id: + await guild.unban(ban_entry.user) + + except (discord.Forbidden, discord.HTTPException) as e: + logger.error( + f"Failed to unban user with ID {case.case_user_id} | Case number {case.case_number} Issue: Failed to unban user. {e}", + ) + return async def setup(bot: Tux) -> None: From c3ca0cc2757824a049d96d64a5637b6ffb4dcbe4 Mon Sep 17 00:00:00 2001 From: petabyte-imo Date: Wed, 4 Sep 2024 14:28:01 +0100 Subject: [PATCH 03/11] Fixed the tempban_check looped task and added some comments --- tux/cogs/moderation/tempban.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index 5940b793e..4b6f90768 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -87,7 +87,6 @@ async def tempban_check(self) -> None: tempbans = [await self.db.case.get_all_cases_by_type(guild.guild_id, CaseType.TEMPBAN) for guild in guilds] # Here, we have 3 nested for loops because for some odd reason, tempbans is a list of a list of lists, very confusing ikr - for tempban in tempbans: for cases in tempbans: for case in cases: From 49d6e45b77761bc56ea21efe22a354146f7c9f97 Mon Sep 17 00:00:00 2001 From: petabyte-imo Date: Wed, 4 Sep 2024 14:31:36 +0100 Subject: [PATCH 04/11] oopsie --- tux/cogs/moderation/tempban.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index 4b6f90768..d07a809bd 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -54,7 +54,7 @@ async def tempban( return moderator = ctx.author - duration = parse_time_string(f"{flags.expires_at}m") + duration = parse_time_string(f"{flags.expires_at}d") expires_at = datetime.now(UTC) + duration if not await self.check_conditions(ctx, member, moderator, "temporarily ban"): From 707e1614bfe54deccc269a60adbb581f108e7f2b Mon Sep 17 00:00:00 2001 From: petabyte-imo Date: Wed, 4 Sep 2024 14:56:57 +0100 Subject: [PATCH 05/11] Added reason to unban --- tux/cogs/moderation/tempban.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index d07a809bd..38fdf5868 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -113,7 +113,7 @@ async def tempban_check(self) -> None: guild_bans = guild.bans() async for ban_entry in guild_bans: if ban_entry.user.id == case.case_user_id: - await guild.unban(ban_entry.user) + await guild.unban(ban_entry.user, reason="Tempban expired") except (discord.Forbidden, discord.HTTPException) as e: logger.error( From e9868f888507dc96ac8c72c9a99db765fd5d1b18 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 4 Sep 2024 13:59:49 +0000 Subject: [PATCH 06/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tux/cogs/moderation/tempban.py | 2 -- tux/database/controllers/guild.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index 38fdf5868..1f788a80f 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -90,14 +90,12 @@ async def tempban_check(self) -> None: for tempban in tempbans: for cases in tempbans: for case in cases: - # Check if the case has expired if case.case_expires_at < datetime.now(UTC): # Get the guild, if that doesnt work then fetch it instead guild = self.bot.get_guild(case.guild_id) if guild is None: - try: guild = await self.bot.fetch_guild(case.guild_id) diff --git a/tux/database/controllers/guild.py b/tux/database/controllers/guild.py index 242570eeb..9fdbf9541 100644 --- a/tux/database/controllers/guild.py +++ b/tux/database/controllers/guild.py @@ -16,4 +16,4 @@ async def delete_guild_by_id(self, guild_id: int) -> None: await self.table.delete(where={"guild_id": guild_id}) async def get_all_guilds(self) -> list[Guild]: - return await self.table.find_many() \ No newline at end of file + return await self.table.find_many() From bbd8fb62aa22ff2730897b8224bf94b8e6e41dfe Mon Sep 17 00:00:00 2001 From: petabyte-imo Date: Wed, 4 Sep 2024 17:48:13 +0100 Subject: [PATCH 07/11] Added a set_tempban_expired and get_expired_tempbans function and changed up tempban a bit --- prisma/schema.prisma | 20 ++------ tux/cogs/moderation/tempban.py | 78 ++++++++++++++++---------------- tux/database/controllers/case.py | 41 +++++++++++++++++ 3 files changed, 85 insertions(+), 54 deletions(-) diff --git a/prisma/schema.prisma b/prisma/schema.prisma index 7fd1a6db8..0eece6470 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -11,14 +11,13 @@ datasource db { } model Guild { - guild_id BigInt @id + guild_id BigInt @id @unique guild_joined_at DateTime? @default(now()) cases Case[] - snippets Snippet[] + guild_config GuildConfig? notes Note[] reminders Reminder[] - guild_config GuildConfig[] - AFK AFKModel[] + snippets Snippet[] Starboard Starboard? StarboardMessage StarboardMessage[] case_count BigInt @default(0) @@ -65,6 +64,7 @@ model Case { case_number BigInt? case_created_at DateTime? @default(now()) case_expires_at DateTime? + case_tempban_expired Boolean? @default(false) guild_id BigInt guild Guild @relation(fields: [guild_id], references: [guild_id]) @@ -115,18 +115,6 @@ model Reminder { @@index([reminder_id, guild_id]) } -model AFKModel { - member_id BigInt @id - nickname String - reason String - since DateTime @default(now()) - guild_id BigInt - guild Guild @relation(fields: [guild_id], references: [guild_id]) - - @@unique([member_id, guild_id]) - @@index([member_id]) -} - model Starboard { guild_id BigInt @id @unique starboard_channel_id BigInt diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index 1f788a80f..b6009cab1 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -54,7 +54,7 @@ async def tempban( return moderator = ctx.author - duration = parse_time_string(f"{flags.expires_at}d") + duration = parse_time_string(f"{flags.expires_at}m") expires_at = datetime.now(UTC) + duration if not await self.check_conditions(ctx, member, moderator, "temporarily ban"): @@ -76,48 +76,50 @@ async def tempban( case_reason=flags.reason, guild_id=ctx.guild.id, case_expires_at=expires_at, + case_tempban_expired=False ) await self.handle_case_response(ctx, CaseType.TEMPBAN, case.case_number, flags.reason, member) - @tasks.loop(minutes=30) + @tasks.loop(seconds=30) async def tempban_check(self) -> None: - # Fetch all guilds and fetch all tempbans for each guild's ID - guilds = await self.db.guild.get_all_guilds() - tempbans = [await self.db.case.get_all_cases_by_type(guild.guild_id, CaseType.TEMPBAN) for guild in guilds] - # Here, we have 3 nested for loops because for some odd reason, tempbans is a list of a list of lists, very confusing ikr - - for tempban in tempbans: - for cases in tempbans: - for case in cases: - # Check if the case has expired - if case.case_expires_at < datetime.now(UTC): - # Get the guild, if that doesnt work then fetch it instead - - guild = self.bot.get_guild(case.guild_id) - if guild is None: - try: - guild = await self.bot.fetch_guild(case.guild_id) - - except (discord.Forbidden, discord.HTTPException) as e: - logger.error( - f"Failed to unban user with ID {case.case_user_id} | Case number {case.case_number} | Issue: Failed to get guild with ID {case.guild_id}. {e}", - ) - return - else: - try: - # Unban the user - - guild_bans = guild.bans() - async for ban_entry in guild_bans: - if ban_entry.user.id == case.case_user_id: - await guild.unban(ban_entry.user, reason="Tempban expired") - - except (discord.Forbidden, discord.HTTPException) as e: - logger.error( - f"Failed to unban user with ID {case.case_user_id} | Case number {case.case_number} Issue: Failed to unban user. {e}", - ) - return + # Get all expired tempbans + expired_temp_bans = await self.db.case.get_expired_tempbans() + logger.debug(f"Checking {len(expired_temp_bans)} expired tempbans. {expired_temp_bans}") + for temp_ban in expired_temp_bans: + #Debug Print + logger.debug(f"Unbanning user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number}") + guild = self.bot.get_guild(temp_ban.guild_id) + + if guild is None: + #Debug Print + logger.debug(f"Fetching guild with ID {temp_ban.guild_id}") + try: + guild = await self.bot.fetch_guild(temp_ban.guild_id) + + except (discord.Forbidden, discord.HTTPException) as e: + logger.error( + f"Failed to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} | Issue: Failed to get guild with ID {temp_ban.guild_id}. {e}", + ) + return + else: + logger.debug(f"Found guild with ID {temp_ban.guild_id}") + try: + # Unban the user + + guild_bans = guild.bans() + async for ban_entry in guild_bans: + if ban_entry.user.id == temp_ban.case_user_id: + await guild.unban(ban_entry.user, reason="Tempban expired") + await self.db.case.set_tempban_expired(temp_ban.case_number, temp_ban.guild_id) + except (discord.Forbidden, discord.HTTPException) as e: + logger.error( + f"Faile+d to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} Issue: Failed to unban user. {e}", + ) + return + #Debug Print + logger.debug(f"Unbanned user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number}") + async def setup(bot: Tux) -> None: diff --git a/tux/database/controllers/case.py b/tux/database/controllers/case.py index b1bd373cd..46fe0281e 100644 --- a/tux/database/controllers/case.py +++ b/tux/database/controllers/case.py @@ -82,6 +82,7 @@ async def insert_case( case_reason: str, case_user_roles: list[int] | None = None, case_expires_at: datetime | None = None, + case_tempban_expired: bool = False, ) -> Case: """ Insert a case into the database. @@ -102,6 +103,8 @@ async def insert_case( The roles of the target of the case. case_expires_at : datetime | None The expiration date of the case. + case_tempban_expired : bool + Whether the tempban has expired (Use only for tempbans). Returns ------- @@ -121,6 +124,7 @@ async def insert_case( "case_reason": case_reason, "case_expires_at": case_expires_at, "case_user_roles": case_user_roles if case_user_roles is not None else [], + "case_tempban_expired": case_tempban_expired, }, ) @@ -307,3 +311,40 @@ async def delete_case_by_number(self, guild_id: int, case_number: int) -> Case | if case is not None: return await self.table.delete(where={"case_id": case.case_id}) return None + + async def get_expired_tempbans(self) -> list[Case]: + """ + Get all cases that have expired tempbans. + + Returns + ------- + list[Case] + A list of cases of the type in the guild. + """ + return await self.table.find_many( + where={ + "case_type": CaseType.TEMPBAN, + "case_expires_at": {"lt": datetime.now()}, + "case_tempban_expired": False + } + ) + async def set_tempban_expired(self, case_number: int, guild_id: int) -> Case | None: + """ + Set a tempban case as expired. + + Parameters + ---------- + case_number : int + The number of the case to delete. + guild_id : int + The ID of the guild to delete the case in. + + Returns + ------- + Case | None + The case if found and deleted, otherwise None. + """ + return await self.table.update( + where={"case_number": case_number, "guild_id": guild_id}, + data={"case_tempban_expired": True} + ) From e293eabbb379e2cfdb70842c2778e84fc7e43274 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 4 Sep 2024 16:49:39 +0000 Subject: [PATCH 08/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tux/cogs/moderation/tempban.py | 13 ++++++------- tux/database/controllers/case.py | 8 ++++---- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index b6009cab1..0f46c6b07 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -76,7 +76,7 @@ async def tempban( case_reason=flags.reason, guild_id=ctx.guild.id, case_expires_at=expires_at, - case_tempban_expired=False + case_tempban_expired=False, ) await self.handle_case_response(ctx, CaseType.TEMPBAN, case.case_number, flags.reason, member) @@ -87,20 +87,20 @@ async def tempban_check(self) -> None: expired_temp_bans = await self.db.case.get_expired_tempbans() logger.debug(f"Checking {len(expired_temp_bans)} expired tempbans. {expired_temp_bans}") for temp_ban in expired_temp_bans: - #Debug Print + # Debug Print logger.debug(f"Unbanning user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number}") guild = self.bot.get_guild(temp_ban.guild_id) if guild is None: - #Debug Print + # Debug Print logger.debug(f"Fetching guild with ID {temp_ban.guild_id}") try: guild = await self.bot.fetch_guild(temp_ban.guild_id) except (discord.Forbidden, discord.HTTPException) as e: logger.error( - f"Failed to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} | Issue: Failed to get guild with ID {temp_ban.guild_id}. {e}", - ) + f"Failed to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} | Issue: Failed to get guild with ID {temp_ban.guild_id}. {e}", + ) return else: logger.debug(f"Found guild with ID {temp_ban.guild_id}") @@ -117,10 +117,9 @@ async def tempban_check(self) -> None: f"Faile+d to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} Issue: Failed to unban user. {e}", ) return - #Debug Print + # Debug Print logger.debug(f"Unbanned user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number}") - async def setup(bot: Tux) -> None: await bot.add_cog(TempBan(bot)) diff --git a/tux/database/controllers/case.py b/tux/database/controllers/case.py index 46fe0281e..db2754177 100644 --- a/tux/database/controllers/case.py +++ b/tux/database/controllers/case.py @@ -325,9 +325,10 @@ async def get_expired_tempbans(self) -> list[Case]: where={ "case_type": CaseType.TEMPBAN, "case_expires_at": {"lt": datetime.now()}, - "case_tempban_expired": False - } + "case_tempban_expired": False, + }, ) + async def set_tempban_expired(self, case_number: int, guild_id: int) -> Case | None: """ Set a tempban case as expired. @@ -345,6 +346,5 @@ async def set_tempban_expired(self, case_number: int, guild_id: int) -> Case | N The case if found and deleted, otherwise None. """ return await self.table.update( - where={"case_number": case_number, "guild_id": guild_id}, - data={"case_tempban_expired": True} + where={"case_number": case_number, "guild_id": guild_id}, data={"case_tempban_expired": True}, ) From c5e0d48b30d3a5e4000c3e925222b9831bc1544a Mon Sep 17 00:00:00 2001 From: petabyte-imo Date: Thu, 5 Sep 2024 14:49:31 +0100 Subject: [PATCH 09/11] Made stuff a bit cleaner and fix some issues related to tempban --- prisma/schema.prisma | 44 +++++++++++++++++++++----------- tux/cogs/moderation/tempban.py | 35 +++++++++++++++---------- tux/database/controllers/case.py | 15 ++++++----- 3 files changed, 58 insertions(+), 36 deletions(-) diff --git a/prisma/schema.prisma b/prisma/schema.prisma index 0eece6470..11184ea48 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -11,13 +11,14 @@ datasource db { } model Guild { - guild_id BigInt @id @unique + guild_id BigInt @id guild_joined_at DateTime? @default(now()) cases Case[] - guild_config GuildConfig? + snippets Snippet[] notes Note[] reminders Reminder[] - snippets Snippet[] + guild_config GuildConfig[] + AFK AFKModel[] Starboard Starboard? StarboardMessage StarboardMessage[] case_count BigInt @default(0) @@ -54,19 +55,19 @@ model GuildConfig { } model Case { - case_id BigInt @id @default(autoincrement()) - case_status Boolean? @default(true) - case_type CaseType - case_reason String - case_moderator_id BigInt - case_user_id BigInt - case_user_roles BigInt[] @default([]) - case_number BigInt? - case_created_at DateTime? @default(now()) - case_expires_at DateTime? + case_id BigInt @id @default(autoincrement()) + case_status Boolean? @default(true) + case_type CaseType + case_reason String + case_moderator_id BigInt + case_user_id BigInt + case_user_roles BigInt[] @default([]) + case_number BigInt? + case_created_at DateTime? @default(now()) + case_expires_at DateTime? case_tempban_expired Boolean? @default(false) - guild_id BigInt - guild Guild @relation(fields: [guild_id], references: [guild_id]) + guild_id BigInt + guild Guild @relation(fields: [guild_id], references: [guild_id]) @@unique([case_number, guild_id]) @@index([case_number, guild_id]) @@ -115,6 +116,18 @@ model Reminder { @@index([reminder_id, guild_id]) } +model AFKModel { + member_id BigInt @id + nickname String + reason String + since DateTime @default(now()) + guild_id BigInt + guild Guild @relation(fields: [guild_id], references: [guild_id]) + + @@unique([member_id, guild_id]) + @@index([member_id]) +} + model Starboard { guild_id BigInt @id @unique starboard_channel_id BigInt @@ -154,4 +167,5 @@ enum CaseType { JAIL UNJAIL SNIPPETUNBAN + UNTEMPBAN } diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index 0f46c6b07..3a15762a0 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -81,44 +81,51 @@ async def tempban( await self.handle_case_response(ctx, CaseType.TEMPBAN, case.case_number, flags.reason, member) - @tasks.loop(seconds=30) + @tasks.loop(hours=1) async def tempban_check(self) -> None: # Get all expired tempbans expired_temp_bans = await self.db.case.get_expired_tempbans() - logger.debug(f"Checking {len(expired_temp_bans)} expired tempbans. {expired_temp_bans}") + for temp_ban in expired_temp_bans: - # Debug Print - logger.debug(f"Unbanning user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number}") - guild = self.bot.get_guild(temp_ban.guild_id) + # This make sure that if get_guild doesnt work, we can instead fetch it from the api + guild = self.bot.get_guild(temp_ban.guild_id) if guild is None: - # Debug Print - logger.debug(f"Fetching guild with ID {temp_ban.guild_id}") try: guild = await self.bot.fetch_guild(temp_ban.guild_id) - except (discord.Forbidden, discord.HTTPException) as e: logger.error( f"Failed to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} | Issue: Failed to get guild with ID {temp_ban.guild_id}. {e}", ) return else: - logger.debug(f"Found guild with ID {temp_ban.guild_id}") try: - # Unban the user + # Unban the user by fetching all the bans and if they exist, unban them guild_bans = guild.bans() async for ban_entry in guild_bans: if ban_entry.user.id == temp_ban.case_user_id: await guild.unban(ban_entry.user, reason="Tempban expired") + + # This just avoids some linting errors + assert isinstance(temp_ban.case_number, int) + + # If the unban goes through, set the tempban as expired and create an untempban case await self.db.case.set_tempban_expired(temp_ban.case_number, temp_ban.guild_id) - except (discord.Forbidden, discord.HTTPException) as e: + await self.db.case.insert_case( + guild_id=temp_ban.guild_id, + case_user_id=temp_ban.case_user_id, + case_moderator_id=temp_ban.case_moderator_id, + case_type=CaseType.UNTEMPBAN, + case_reason="Expired tempban", + case_tempban_expired=True, + ) + + except (discord.Forbidden, discord.HTTPException, discord.NotFound) as e: logger.error( - f"Faile+d to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} Issue: Failed to unban user. {e}", + f"Failed to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} Issue: Failed to unban user. {e}", ) return - # Debug Print - logger.debug(f"Unbanned user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number}") async def setup(bot: Tux) -> None: diff --git a/tux/database/controllers/case.py b/tux/database/controllers/case.py index db2754177..063cecce1 100644 --- a/tux/database/controllers/case.py +++ b/tux/database/controllers/case.py @@ -1,4 +1,4 @@ -from datetime import datetime +from datetime import UTC, datetime from prisma.enums import CaseType from prisma.models import Case, Guild @@ -324,12 +324,12 @@ async def get_expired_tempbans(self) -> list[Case]: return await self.table.find_many( where={ "case_type": CaseType.TEMPBAN, - "case_expires_at": {"lt": datetime.now()}, + "case_expires_at": {"lt": datetime.now(UTC)}, "case_tempban_expired": False, }, ) - async def set_tempban_expired(self, case_number: int, guild_id: int) -> Case | None: + async def set_tempban_expired(self, case_number: int, guild_id: int) -> int: """ Set a tempban case as expired. @@ -342,9 +342,10 @@ async def set_tempban_expired(self, case_number: int, guild_id: int) -> Case | N Returns ------- - Case | None - The case if found and deleted, otherwise None. + int + The total number of Case records that were updated. """ - return await self.table.update( - where={"case_number": case_number, "guild_id": guild_id}, data={"case_tempban_expired": True}, + return await self.table.update_many( + where={"case_number": case_number, "guild_id": guild_id}, + data={"case_tempban_expired": True}, ) From 0f2a588dfa10c3d37b3fafccd36cde543ad37057 Mon Sep 17 00:00:00 2001 From: petabyte-imo Date: Thu, 5 Sep 2024 14:54:13 +0100 Subject: [PATCH 10/11] oopsiess --- tux/cogs/moderation/tempban.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index 3a15762a0..547064a20 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -54,7 +54,7 @@ async def tempban( return moderator = ctx.author - duration = parse_time_string(f"{flags.expires_at}m") + duration = parse_time_string(f"{flags.expires_at}d") expires_at = datetime.now(UTC) + duration if not await self.check_conditions(ctx, member, moderator, "temporarily ban"): From d564cad580c1f67e40116422c73e329836ef96e1 Mon Sep 17 00:00:00 2001 From: wlinator Date: Thu, 5 Sep 2024 10:29:55 -0400 Subject: [PATCH 11/11] Refactor tempban & optimize --- tux/cogs/moderation/tempban.py | 79 ++++++++++++-------------------- tux/database/controllers/case.py | 24 +++++++--- tux/utils/flags.py | 4 +- 3 files changed, 50 insertions(+), 57 deletions(-) diff --git a/tux/cogs/moderation/tempban.py b/tux/cogs/moderation/tempban.py index 547064a20..d726fc90f 100644 --- a/tux/cogs/moderation/tempban.py +++ b/tux/cogs/moderation/tempban.py @@ -30,7 +30,7 @@ async def tempban( flags: TempBanFlags, ) -> None: """ - Ban a member from the server. + Temporarily ban a member from the server. Parameters ---------- @@ -39,7 +39,7 @@ async def tempban( member : discord.Member The member to ban. flags : TempBanFlags - The flags for the command. (reason: str, purge_days: int, silent: bool, expires_at: int The number of days the ban will last for.) + The flags for the command. Raises ------ @@ -48,22 +48,19 @@ async def tempban( discord.HTTPException If an error occurs while banning the user. """ - - if ctx.guild is None: + if not ctx.guild: logger.warning("Ban command used outside of a guild context.") return - moderator = ctx.author + if not await self.check_conditions(ctx, member, ctx.author, "temporarily ban"): + return + duration = parse_time_string(f"{flags.expires_at}d") expires_at = datetime.now(UTC) + duration - if not await self.check_conditions(ctx, member, moderator, "temporarily ban"): - return - try: await self.send_dm(ctx, flags.silent, member, flags.reason, action="temporarily banned") await ctx.guild.ban(member, reason=flags.reason, delete_message_days=flags.purge_days) - except (discord.Forbidden, discord.HTTPException) as e: logger.error(f"Failed to temporarily ban {member}. {e}") await ctx.send(f"Failed to temporarily ban {member}. {e}", delete_after=30, ephemeral=True) @@ -83,49 +80,33 @@ async def tempban( @tasks.loop(hours=1) async def tempban_check(self) -> None: - # Get all expired tempbans expired_temp_bans = await self.db.case.get_expired_tempbans() for temp_ban in expired_temp_bans: - # This make sure that if get_guild doesnt work, we can instead fetch it from the api - - guild = self.bot.get_guild(temp_ban.guild_id) - if guild is None: - try: - guild = await self.bot.fetch_guild(temp_ban.guild_id) - except (discord.Forbidden, discord.HTTPException) as e: - logger.error( - f"Failed to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} | Issue: Failed to get guild with ID {temp_ban.guild_id}. {e}", - ) - return - else: - try: - # Unban the user by fetching all the bans and if they exist, unban them - - guild_bans = guild.bans() - async for ban_entry in guild_bans: - if ban_entry.user.id == temp_ban.case_user_id: - await guild.unban(ban_entry.user, reason="Tempban expired") - - # This just avoids some linting errors - assert isinstance(temp_ban.case_number, int) - - # If the unban goes through, set the tempban as expired and create an untempban case - await self.db.case.set_tempban_expired(temp_ban.case_number, temp_ban.guild_id) - await self.db.case.insert_case( - guild_id=temp_ban.guild_id, - case_user_id=temp_ban.case_user_id, - case_moderator_id=temp_ban.case_moderator_id, - case_type=CaseType.UNTEMPBAN, - case_reason="Expired tempban", - case_tempban_expired=True, - ) - - except (discord.Forbidden, discord.HTTPException, discord.NotFound) as e: - logger.error( - f"Failed to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number} Issue: Failed to unban user. {e}", - ) - return + guild = self.bot.get_guild(temp_ban.guild_id) or await self.bot.fetch_guild(temp_ban.guild_id) + if not guild: + logger.error(f"Failed to get guild with ID {temp_ban.guild_id} for tempban check.") + continue + + try: + ban_entry = await guild.fetch_ban(discord.Object(id=temp_ban.case_user_id)) + await guild.unban(ban_entry.user, reason=f"Tempban expired | Case number: {temp_ban.case_number}") + + await self.db.case.set_tempban_expired(temp_ban.case_number, temp_ban.guild_id) + await self.db.case.insert_case( + guild_id=temp_ban.guild_id, + case_user_id=temp_ban.case_user_id, + case_moderator_id=temp_ban.case_moderator_id, + case_type=CaseType.UNTEMPBAN, + case_reason="Expired tempban", + case_tempban_expired=True, + ) + logger.debug(f"Unbanned user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number}") + + except (discord.Forbidden, discord.HTTPException, discord.NotFound) as e: + logger.error( + f"Failed to unban user with ID {temp_ban.case_user_id} | Case number {temp_ban.case_number}. Error: {e}", + ) async def setup(bot: Tux) -> None: diff --git a/tux/database/controllers/case.py b/tux/database/controllers/case.py index 063cecce1..ae02373d8 100644 --- a/tux/database/controllers/case.py +++ b/tux/database/controllers/case.py @@ -329,23 +329,35 @@ async def get_expired_tempbans(self) -> list[Case]: }, ) - async def set_tempban_expired(self, case_number: int, guild_id: int) -> int: + async def set_tempban_expired(self, case_number: int | None, guild_id: int) -> int | None: """ Set a tempban case as expired. Parameters ---------- case_number : int - The number of the case to delete. + The number of the case to update. guild_id : int - The ID of the guild to delete the case in. + The ID of the guild the case belongs to. Returns ------- - int - The total number of Case records that were updated. + Optional[int] + The number of Case records updated (1) if successful, None if no records were found, + or raises an exception if multiple records were affected. """ - return await self.table.update_many( + if case_number is None: + msg = "Case number not found" + raise ValueError(msg) + + result = await self.table.update_many( where={"case_number": case_number, "guild_id": guild_id}, data={"case_tempban_expired": True}, ) + if result == 1: + return result + if result == 0: + return None + + msg = f"Multiple records ({result}) were affected when updating case {case_number} in guild {guild_id}" + raise ValueError(msg) diff --git a/tux/utils/flags.py b/tux/utils/flags.py index 6869ea098..7fbfb3480 100644 --- a/tux/utils/flags.py +++ b/tux/utils/flags.py @@ -91,9 +91,9 @@ class TempBanFlags(commands.FlagConverter, case_insensitive=True, delimiter=" ", default=MISSING, ) expires_at: int = commands.flag( - name="expires_at", + name="duration", description="Number of days the ban will last for.", - aliases=["t", "d", "e", "duration", "expires", "time"], + aliases=["t", "d", "e", "expires", "time"], ) purge_days: int = commands.flag( name="purge_days",