Merge pull request 'Finish the addrole/removerole commands and surrounding functionality' (#24) from addrole into main

Reviewed-on: SeaswimmerTheFsh/SeaCogs#24

aurora/aurora.py

@@ -13,6 +13,7 @@ from datetime import datetime, timedelta, timezone
 from math import ceil
 
 import discord
+from discord import Object
 from discord.ext import tasks
 from redbot.core import app_commands, commands, data_manager
 from redbot.core.app_commands import Choice
@@ -30,7 +31,7 @@ from aurora.utilities.config import config, register_config
 from aurora.utilities.database import connect, create_guild_table, fetch_case, mysql_log
 from aurora.utilities.factory import addrole_embed, case_factory, changes_factory, evidenceformat_factory, guild_embed, immune_embed, message_factory, overrides_embed
 from aurora.utilities.logger import logger
-from aurora.utilities.utils import check_moddable, check_permissions, convert_timedelta_to_str, fetch_channel_dict, fetch_user_dict, generate_dict, log, send_evidenceformat, timedelta_from_relativedelta
+from aurora.utilities.utils import check_moddable, check_permissions, convert_timedelta_to_str, fetch_channel_dict, fetch_user_dict, generate_dict, get_footer_image, log, send_evidenceformat, timedelta_from_relativedelta
 
 
 class Aurora(commands.Cog):
@@ -113,6 +114,20 @@ class Aurora(commands.Cog):
             except ConnectionRefusedError:
                 return
 
+    @commands.Cog.listener("on_member_join")
+    async def addrole_on_member_join(self, member: discord.Member):
+        """This method automatically adds roles to users when they join the server."""
+        if not await self.bot.cog_disabled_in_guild(self, member.guild):
+            query = f"""SELECT moderation_id, role_id, reason FROM moderation_{member.guild.id} WHERE target_id = ? AND moderation_type = 'ADDROLE' AND expired = 0 AND resolved = 0;"""
+            database = connect()
+            cursor = database.cursor()
+            cursor.execute(query, (member.id,))
+            results = cursor.fetchall()
+            for result in results:
+                role = member.guild.get_role(result[1])
+                reason = result[2]
+                await member.add_roles(role, reason=f"Role automatically added on member rejoin for: {reason} (Case #{result[0]:,})")
+
     @commands.Cog.listener("on_audit_log_entry_create")
     async def autologger(self, entry: discord.AuditLogEntry):
         """This method automatically logs moderations done by users manually ("right clicks")."""
@@ -209,7 +224,7 @@ class Aurora(commands.Cog):
                     moderation_type="note",
                     response=await interaction.original_response(),
                 )
-                await target.send(embed=embed)
+                await target.send(embed=embed, file=get_footer_image(self))
             except discord.errors.HTTPException:
                 pass
 
@@ -268,7 +283,7 @@ class Aurora(commands.Cog):
                     moderation_type="warned",
                     response=await interaction.original_response(),
                 )
-                await target.send(embed=embed)
+                await target.send(embed=embed, file=get_footer_image(self))
             except discord.errors.HTTPException:
                 pass
 
@@ -366,16 +381,16 @@ class Aurora(commands.Cog):
                     duration=parsed_time,
                     role=role,
                 )
-                await target.send(embed=embed)
+                await target.send(embed=embed, file=get_footer_image(self))
             except discord.errors.HTTPException:
                 pass
 
         await target.add_roles(
             role,
-            reason=f"Role added by {interaction.user.id}{(' for ' + {humanize_timedelta(timedelta=parsed_time)} if parsed_time != 'NULL' else '')} for: {reason}",
+            reason=f"Role added by {interaction.user.id}{' for ' + humanize_timedelta(timedelta=parsed_time) if parsed_time != 'NULL' else ''} for: {reason}",
         )
         response: discord.WebhookMessage = await interaction.followup.send(
-            content=f"{target.mention} has been given the {role.mention} role{(' for ' + {humanize_timedelta(timedelta=parsed_time)} if parsed_time != 'NULL' else '')}!\n**Reason** - `{reason}`"
+            content=f"{target.mention} has been given the {role.mention} role{' for ' + humanize_timedelta(timedelta=parsed_time) if parsed_time != 'NULL' else ''}!\n**Reason** - `{reason}`"
         )
 
         moderation_id = await mysql_log(
@@ -389,7 +404,113 @@ class Aurora(commands.Cog):
             reason,
         )
         await response.edit(
-            content=f"{target.mention} has been given the {role.mention} role{(' for ' + {humanize_timedelta(timedelta=parsed_time)} if parsed_time != 'NULL' else '')}! (Case `#{moderation_id:,}`)\n**Reason** - `{reason}`",
+            content=f"{target.mention} has been given the {role.mention} role{' for ' + humanize_timedelta(timedelta=parsed_time) if parsed_time != 'NULL' else ''}! (Case `#{moderation_id:,}`)\n**Reason** - `{reason}`",
+        )
+        await log(interaction, moderation_id)
+
+        case = await fetch_case(moderation_id, interaction.guild.id)
+        await send_evidenceformat(interaction, case)
+
+    @app_commands.command(name="removerole")
+    async def removerole(
+        self,
+        interaction: discord.Interaction,
+        target: discord.Member,
+        role: discord.Role,
+        reason: str,
+        duration: str = None,
+        silent: bool = None,
+    ):
+        """Remove a role from a user.
+
+        Parameters
+        -----------
+        target: discord.Member
+            Who are you removing a role from?
+        role: discord.Role
+            What role are you removing from the target?
+        reason: str
+            Why are you removing a role from this user?
+        duration: str
+            How long are you removing this role for?
+        silent: bool
+            Should the user be messaged?"""
+        addrole_whitelist = await config.guild(interaction.guild).addrole_whitelist()
+
+        if not addrole_whitelist:
+            await interaction.response.send_message(
+                content=error("There are no whitelisted roles set for this server!"),
+                ephemeral=True,
+            )
+            return
+
+        if duration is not None:
+            parsed_time = parse_timedelta(duration)
+            if parsed_time is None:
+                await interaction.response.send_message(
+                    content=error("Please provide a valid duration!"), ephemeral=True
+                )
+                return
+        else:
+            parsed_time = "NULL"
+
+        if role.id not in addrole_whitelist:
+            await interaction.response.send_message(
+                content=error("That role isn't whitelisted!"), ephemeral=True
+            )
+            return
+
+        if not await check_moddable(
+            target, interaction, ["moderate_members", "manage_roles"]
+        ):
+            return
+
+        if role.id not in [user_role.id for user_role in target.roles]:
+            await interaction.response.send_message(
+                content=error(f"{target.mention} does not have this role!"),
+                ephemeral=True,
+            )
+            return
+
+        await interaction.response.defer()
+        if silent is None:
+            silent = not await config.guild(interaction.guild).dm_users()
+        if silent is False:
+            try:
+                embed = await message_factory(
+                    await self.bot.get_embed_color(interaction.channel),
+                    guild=interaction.guild,
+                    moderator=interaction.user,
+                    reason=reason,
+                    moderation_type="removerole",
+                    response=await interaction.original_response(),
+                    duration=parsed_time,
+                    role=role,
+                )
+                await target.send(embed=embed, file=get_footer_image(self))
+            except discord.errors.HTTPException:
+                pass
+
+        await target.remove_roles(
+            role,
+            reason=f"Role removed by {interaction.user.id}{' for ' + humanize_timedelta(timedelta=parsed_time) if parsed_time != 'NULL' else ''} for: {reason}",
+        )
+        response: discord.WebhookMessage = await interaction.followup.send(
+            content=f"{target.mention} has had the {role.mention} role removed{' for ' + humanize_timedelta(timedelta=parsed_time) if parsed_time != 'NULL' else ''}!\n**Reason** - `{reason}`"
+        )
+
+        moderation_id = await mysql_log(
+            interaction.guild.id,
+            interaction.user.id,
+            "REMOVEROLE",
+            "USER",
+            target.id,
+            role.id,
+            parsed_time,
+            reason,
+        )
+        await response.edit(
+            content=f"{target.mention} has had the {role.mention} role removed{' for ' + humanize_timedelta(timedelta=parsed_time) if parsed_time != 'NULL' else ''}! (Case `#{moderation_id:,}`)\n**Reason** - `{reason}`",
         )
         await log(interaction, moderation_id)
 
@@ -462,7 +583,7 @@ class Aurora(commands.Cog):
                     response=await interaction.original_response(),
                     duration=parsed_time,
                 )
-                await target.send(embed=embed)
+                await target.send(embed=embed, file=get_footer_image(self))
             except discord.errors.HTTPException:
                 pass
 
@@ -537,7 +658,7 @@ class Aurora(commands.Cog):
                     moderation_type="unmuted",
                     response=await interaction.original_response(),
                 )
-                await target.send(embed=embed)
+                await target.send(embed=embed, file=get_footer_image(self))
             except discord.errors.HTTPException:
                 pass
 
@@ -596,7 +717,7 @@ class Aurora(commands.Cog):
                     moderation_type="kicked",
                     response=await interaction.original_response(),
                 )
-                await target.send(embed=embed)
+                await target.send(embed=embed, file=get_footer_image(self))
             except discord.errors.HTTPException:
                 pass
 
@@ -700,7 +821,7 @@ class Aurora(commands.Cog):
                     response=await interaction.original_response(),
                     duration=parsed_time,
                 )
-                await target.send(embed=embed)
+                await target.send(embed=embed, file=get_footer_image(self))
             except discord.errors.HTTPException:
                 pass
 
@@ -744,7 +865,7 @@ class Aurora(commands.Cog):
                         moderation_type="banned",
                         response=await interaction.original_response(),
                     )
-                    await target.send(embed=embed)
+                    await target.send(embed=embed, file=get_footer_image(self))
                 except discord.errors.HTTPException:
                     pass
 
@@ -827,7 +948,7 @@ class Aurora(commands.Cog):
                     moderation_type="unbanned",
                     response=await interaction.original_response(),
                 )
-                await target.send(embed=embed)
+                await target.send(embed=embed, file=get_footer_image(self))
             except discord.errors.HTTPException:
                 pass
 
@@ -1306,7 +1427,7 @@ class Aurora(commands.Cog):
                         os.remove(filename)
                         return
                     await interaction.response.send_message(
-                        content=box({json.dumps(case_dict, indent=2)}),
+                        content=box(json.dumps(case_dict, indent=2), 'json'),
                         ephemeral=ephemeral,
                     )
                     return
@@ -1490,7 +1611,9 @@ class Aurora(commands.Cog):
         current_time = time.time()
         database = connect()
         cursor = database.cursor()
-        global_num = 0
+        global_unban_num = 0
+        global_addrole_num = 0
+        global_removerole_num = 0
 
         guilds: list[discord.Guild] = self.bot.guilds
         for guild in guilds:
@@ -1508,7 +1631,7 @@ class Aurora(commands.Cog):
                 target_ids = [row[0] for row in result]
                 moderation_ids = [row[1] for row in result]
 
-                num = 0
+                unban_num = 0
                 for target_id, moderation_id in zip(target_ids, moderation_ids):
                     user: discord.User = await self.bot.fetch_user(target_id)
                     name = (
@@ -1529,7 +1652,7 @@ class Aurora(commands.Cog):
                         )
 
                         try:
-                            await user.send(embed=embed)
+                            await user.send(embed=embed, file=get_footer_image(self))
                         except discord.errors.HTTPException:
                             pass
 
@@ -1540,7 +1663,7 @@ class Aurora(commands.Cog):
                             guild.name,
                             guild.id,
                         )
-                        num = num + 1
+                        unban_num = unban_num + 1
                     except (
                         discord.errors.NotFound,
                         discord.errors.Forbidden,
@@ -1555,12 +1678,10 @@ 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')"
+                removerole_num = 0
-                cursor.execute(expiry_query, (time.time(),))
+                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"
-
-                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"
                 try:
-                    cursor.execute(blacklist_query, (time.time(),))
+                    cursor.execute(addrole_query, (time.time(),))
                     result = cursor.fetchall()
                 except sqlite3.OperationalError:
                     continue
@@ -1572,27 +1693,72 @@ 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)
+                        await member.remove_roles(
-                        if role is None:
+                            Object(role_id), reason=f"Automatic role removal from case #{moderation_id}"
-                            raise discord.errors.NotFound
+                        )
+
+                        removerole_num = removerole_num + 1
                     except (
                         discord.errors.NotFound,
                         discord.errors.Forbidden,
                         discord.errors.HTTPException,
-                    ):
+                    ) as e:
+                        logger.error(
+                            "Removing the role %s from user %s failed due to: \n%s",
+                            role_id,
+                            target_id,
+                            e,
+                        )
                         continue
 
+                addrole_num = 0
+                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)
+
+                        await member.add_roles(
+                            Object(role_id), reason=f"Automatic role addition from case #{moderation_id}"
+                        )
+
+                        addrole_num = addrole_num + 1
+                    except (
+                        discord.errors.NotFound,
+                        discord.errors.Forbidden,
+                        discord.errors.HTTPException,
+                    ) as e:
+                        logger.error("Adding the role %s to user %s failed due to: \n%s", role_id, target_id, e)
+                        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",
+                    "Completed expiry loop for %s (%s) in %sms with %s users unbanned, %s roles added, and %s roles removed",
                     guild.name,
                     guild.id,
                     f"{per_guild_completion_time:.6f}",
-                    num,
+                    unban_num,
+                    addrole_num,
+                    removerole_num,
                 )
-                global_num = global_num + num
+                global_unban_num = global_unban_num + unban_num
+                global_addrole_num = global_addrole_num + addrole_num
+                global_removerole_num = global_removerole_num + removerole_num
 
         database.commit()
         cursor.close()
@@ -1600,9 +1766,11 @@ class Aurora(commands.Cog):
 
         completion_time = (time.time() - current_time) * 1000
         logger.debug(
-            "Completed expiry loop in %sms with %s users unbanned",
+            "Completed expiry loop in %sms with %s users unbanned, %s roles added, and %s roles removed",
             f"{completion_time:.6f}",
-            global_num,
+            global_unban_num,
+            global_addrole_num,
+            global_removerole_num,
         )
 
     ########################################################################################################################

aurora/utilities/factory.py

@@ -2,16 +2,12 @@
 from datetime import datetime, timedelta
 from typing import Union
 
-from discord import (Color, Embed, Guild, Interaction, InteractionMessage,
+from discord import Color, Embed, Guild, Interaction, InteractionMessage, Member, Role, User
-                     Member, Role, User)
 from redbot.core import commands
-from redbot.core.utils.chat_formatting import (bold, box, error,
+from redbot.core.utils.chat_formatting import bold, box, error, humanize_timedelta, warning
-                                               humanize_timedelta, warning)
 
 from aurora.utilities.config import config
-from aurora.utilities.utils import (fetch_channel_dict, fetch_user_dict,
+from aurora.utilities.utils import fetch_channel_dict, fetch_user_dict, get_bool_emoji, get_next_case_number, get_pagesize_str
-                                    get_bool_emoji, get_next_case_number,
-                                    get_pagesize_str)
 
 
 async def message_factory(
@@ -50,6 +46,8 @@ async def message_factory(
     else:
         guild_name = guild.name
 
+    title = moderation_type
+
     if moderation_type in ["tempbanned", "muted"] and duration:
         embed_duration = f" for {humanize_timedelta(timedelta=duration)}"
     else:
@@ -59,13 +57,17 @@ async def message_factory(
         embed_desc = "received a"
     elif moderation_type == "addrole":
         embed_desc = f"received the {role.name} role"
+        title = "Role Added"
+        moderation_type = ""
     elif moderation_type == "removerole":
         embed_desc = f"lost the {role.name} role"
+        title = "Role Removed"
+        moderation_type = ""
     else:
         embed_desc = "been"
 
     embed = Embed(
-        title=str.title(moderation_type),
+        title=str.title(title),
         description=f"You have {embed_desc} {moderation_type}{embed_duration} in {guild_name}.",
         color=color,
         timestamp=datetime.now(),
@@ -85,7 +87,7 @@ async def message_factory(
 
     embed.set_footer(
         text=f"Case #{await get_next_case_number(guild.id):,}",
-        icon_url="https://cdn.discordapp.com/attachments/1070822161389994054/1159469476773904414/arrow-right-circle-icon-512x512-2p1e2aaw.png?ex=65312319&is=651eae19&hm=3cebdd28e805c13a79ec48ef87c32ca532ffa6b9ede2e48d0cf8e5e81f3a6818&",
+        icon_url="attachment://arrow.png",
     )
 
     return embed
@@ -267,6 +269,9 @@ async def case_factory(interaction: Interaction, case_dict: dict) -> Embed:
         else "\n**Changes:** 0"
     )
 
+    if case_dict["role_id"]:
+        embed.description += f"\n**Role:** <@&{case_dict['role_id']}>"
+
     if case_dict["metadata"]:
         if case_dict["metadata"]["imported_from"]:
             embed.description += (

aurora/utilities/utils.py

@@ -5,9 +5,9 @@ from datetime import timedelta as td
 from typing import Optional, Union
 
 from dateutil.relativedelta import relativedelta as rd
-from discord import Guild, Interaction, Member, SelectOption, User
+from discord import File, Guild, Interaction, Member, SelectOption, User
 from discord.errors import Forbidden, NotFound
-from redbot.core import commands
+from redbot.core import commands, data_manager
 from redbot.core.utils.chat_formatting import error
 
 from .config import config
@@ -291,3 +291,8 @@ def timedelta_from_relativedelta(relativedelta: rd) -> td:
     now = datetime.now()
     then = now - relativedelta
     return now - then
+
+def get_footer_image(coginstance: commands.Cog) -> File:
+    """Returns the footer image for the embeds."""
+    image_path = data_manager.bundled_data_path(coginstance) / "arrow.png"
+    return File(image_path, filename="arrow.png", description="arrow")