From 786bc95eeb05760e5a4d22e79ba6322ae78ea374 Mon Sep 17 00:00:00 2001
From: profitroll <vozhd.kk@gmail.com>
Date: Mon, 8 May 2023 15:45:00 +0200
Subject: [PATCH] Changed admin check

---
 classes/holo_user.py | 40 ++++++++++++++++++++++++++++++++++++++++
 cogs/data.py         | 11 +++++++++--
 config_example.json  |  1 +
 3 files changed, 50 insertions(+), 2 deletions(-)

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": {},