From f6b827c64fce78cc0f623e81f74ea7091687b734 Mon Sep 17 00:00:00 2001
From: SeaswimmerTheFsh <seaswimmerthefsh@gmail.com>
Date: Fri, 3 May 2024 20:09:48 -0400
Subject: [PATCH] feat(aurora): added expiry handling for addrole and
 removerole

---
 aurora/aurora.py | 48 +++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 41 insertions(+), 7 deletions(-)

diff --git a/aurora/aurora.py b/aurora/aurora.py
index 42804a7..144b2a8 100644
--- a/aurora/aurora.py
+++ b/aurora/aurora.py
@@ -1670,12 +1670,9 @@ class Aurora(commands.Cog):
                             e,
                         )
 
-                expiry_query = f"UPDATE `moderation_{guild.id}` SET expired = 1 WHERE (end_timestamp != 0 AND end_timestamp <= ? AND expired = 0 AND moderation_type != 'BLACKLIST') OR (expired = 0 AND resolved = 1  AND moderation_type != 'BLACKLIST')"
-                cursor.execute(expiry_query, (time.time(),))
-
-                blacklist_query = f"SELECT target_id, moderation_id, role_id FROM moderation_{guild.id} WHERE end_timestamp != 0 AND end_timestamp <= ? AND moderation_type = 'BLACKLIST' AND expired = 0"
+                addrole_query = f"SELECT target_id, moderation_id, role_id FROM moderation_{guild.id} WHERE end_timestamp != 0 AND end_timestamp <= ? AND moderation_type = 'ADDROLE' AND expired = 0"
                 try:
-                    cursor.execute(blacklist_query, (time.time(),))
+                    cursor.execute(addrole_query, (time.time(),))
                     result = cursor.fetchall()
                 except sqlite3.OperationalError:
                     continue
@@ -1687,11 +1684,15 @@ class Aurora(commands.Cog):
                     target_ids, moderation_ids, role_ids
                 ):
                     try:
-                        # member: discord.Member = await guild.fetch_member(target_id)
+                        member = await guild.fetch_member(target_id)
 
-                        role: discord.Role = guild.get_role(role_id)
+                        role = guild.get_role(role_id)
                         if role is None:
                             raise discord.errors.NotFound
+
+                        await member.remove_roles(
+                            role, reason=f"Automatic role removal from case #{moderation_id}"
+                        )
                     except (
                         discord.errors.NotFound,
                         discord.errors.Forbidden,
@@ -1699,6 +1700,39 @@ class Aurora(commands.Cog):
                     ):
                         continue
 
+                removerole_query = f"SELECT target_id, moderation_id, role_id FROM moderation_{guild.id} WHERE end_timestamp != 0 AND end_timestamp <= ? AND moderation_type = 'REMOVEROLE' AND expired = 0"
+                try:
+                    cursor.execute(removerole_query, (time.time(),))
+                    result = cursor.fetchall()
+                except sqlite3.OperationalError:
+                    continue
+                target_ids = [row[0] for row in result]
+                moderation_ids = [row[1] for row in result]
+                role_ids = [row[2] for row in result]
+
+                for target_id, moderation_id, role_id in zip(
+                    target_ids, moderation_ids, role_ids
+                ):
+                    try:
+                        member = await guild.fetch_member(target_id)
+
+                        role = guild.get_role(role_id)
+                        if role is None:
+                            raise discord.errors.NotFound
+
+                        await member.add_roles(
+                            role, reason=f"Automatic role addition from case #{moderation_id}"
+                        )
+                    except (
+                        discord.errors.NotFound,
+                        discord.errors.Forbidden,
+                        discord.errors.HTTPException,
+                    ):
+                        continue
+
+                expiry_query = f"UPDATE `moderation_{guild.id}` SET expired = 1 WHERE (end_timestamp != 0 AND end_timestamp <= ? AND expired = 0) OR (expired = 0 AND resolved = 1"
+                cursor.execute(expiry_query, (time.time(),))
+
                 per_guild_completion_time = (time.time() - time_per_guild) * 1000
                 logger.debug(
                     "Completed expiry loop for %s (%s) in %sms with %s users unbanned",