diff --git a/classes/holo_user.py b/classes/holo_user.py index 08633b1..4658128 100644 --- a/classes/holo_user.py +++ b/classes/holo_user.py @@ -5,6 +5,7 @@ import discord import discord.member from modules.database import col_users, col_warnings +from modules.utils import config_get logger = logging.getLogger(__name__) @@ -98,6 +99,45 @@ class HoloUser: ) logger.info(f"Set attribute {key} of user {self.id} to {value}") + async def is_moderator( + self, member: Union[discord.User, discord.Member, discord.member.Member] + ) -> bool: + """Check if user is moderator or council member + + ### Args: + * `member` (Union[discord.User, discord.Member, discord.member.Member]): Member object + + ### Returns: + `bool`: `True` if member is a moderator or member of council and `False` if not + """ + if isinstance(member, discord.User): + return False + moderator_role = await config_get("moderator", "roles") + council_role = await config_get("council", "roles") + for role in member.roles: + if role.id == moderator_role or role.id == council_role: + return True + return False + + async def is_council( + self, member: Union[discord.User, discord.Member, discord.member.Member] + ) -> bool: + """Check if user is a member of council + + ### Args: + * `member` (Union[discord.User, discord.Member, discord.member.Member]): Member object + + ### Returns: + `bool`: `True` if member is a member of council and `False` if not + """ + if isinstance(member, discord.User): + return False + council_role = await config_get("council", "roles") + for role in member.roles: + if role.id == council_role: + return True + return False + # def purge(self) -> None: # """Completely remove data from database. Will not remove transactions logs and warnings.""" # col_users.delete_one(filter={"_id": self.db_id}) diff --git a/cogs/data.py b/cogs/data.py index e83685d..7ef39f6 100644 --- a/cogs/data.py +++ b/cogs/data.py @@ -7,6 +7,7 @@ from discord import ApplicationContext, Embed, File, option from discord import utils as ds_utils from discord.commands import SlashCommandGroup from discord.ext import commands +from classes.holo_user import HoloUser from enums.colors import Color from modules.database import col_users @@ -32,7 +33,10 @@ class Data(commands.Cog): ) async def data_export_cmd(self, ctx: ApplicationContext, kind: str): await ctx.defer() - if ctx.user.id in await config_get("admins"): + holo_user = HoloUser(ctx.author) + if (ctx.user.id in await config_get("admins")) or ( + await holo_user.is_council(ctx.author) + ): logging.info( f"Moderator {guild_name(ctx.user)} exported current users list" ) @@ -94,7 +98,10 @@ class Data(commands.Cog): ) async def data_migrate_cmd(self, ctx: ApplicationContext, kind: str): await ctx.defer() - if ctx.user.id in await config_get("admins"): + holo_user = HoloUser(ctx.author) + if (ctx.user.id in await config_get("admins")) or ( + await holo_user.is_council(ctx.author) + ): logging.info( f"Moderator {guild_name(ctx.user)} started migration of all members to the database" ) diff --git a/config_example.json b/config_example.json index 2f6525b..eacf44b 100644 --- a/config_example.json +++ b/config_example.json @@ -36,6 +36,7 @@ "voice": {} }, "roles": { + "council": 0, "moderator": 0 }, "bots": {},