From 640ab65e7a7acce0ac43aa540ad77a6097dc8b5d Mon Sep 17 00:00:00 2001
From: profitroll <vozhd.kk@gmail.com>
Date: Sat, 6 May 2023 18:48:04 +0200
Subject: [PATCH] Improved data usage handling

---
 cogs/data.py | 117 +++++++++++++++++++++++++++++++++++++++++----------
 main.py      |   2 +-
 2 files changed, 96 insertions(+), 23 deletions(-)

diff --git a/cogs/data.py b/cogs/data.py
index 3ad88ea..e83685d 100644
--- a/cogs/data.py
+++ b/cogs/data.py
@@ -1,15 +1,15 @@
 import logging
-from datetime import timedelta
-from os import getpid, makedirs, system
+from os import makedirs
 from pathlib import Path
 from uuid import uuid4
 
-from discord import ApplicationContext, Embed, Option, User, slash_command, File
+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
 from modules.utils import config_get
 from modules.utils_sync import config_get_sync, guild_name, json_write_sync
 
@@ -20,15 +20,17 @@ class Data(commands.Cog):
     def __init__(self, client):
         self.client = client
 
-    @slash_command(
+    data = SlashCommandGroup("data", "Керування даними користувачів")
+
+    @data.command(
         name="export",
         description="Експортувати дані",
         guild_ids=[config_get_sync("guild")],
     )
-    async def warn_cmd(
-        self,
-        ctx: ApplicationContext,
-    ):
+    @option(
+        "kind", description="Тип даних, які треба експортувати", choices=["Користувачі"]
+    )
+    async def data_export_cmd(self, ctx: ApplicationContext, kind: str):
         await ctx.defer()
         if ctx.user.id in await config_get("admins"):
             logging.info(
@@ -36,21 +38,25 @@ class Data(commands.Cog):
             )
             makedirs("tmp", exist_ok=True)
             uuid = str(uuid4())
-            users = []
 
-            for member in ctx.guild.members:
-                users.append(
-                    {
-                        "id": member.id,
-                        "nick": member.nick,
-                        "username": f"{member.name}#{member.discriminator}",
-                        "bot": member.bot,
-                    }
+            if kind == "Користувачі":
+                users = []
+
+                for member in ctx.guild.members:
+                    users.append(
+                        {
+                            "id": member.id,
+                            "nick": member.nick,
+                            "username": f"{member.name}#{member.discriminator}",
+                            "bot": member.bot,
+                        }
+                    )
+
+                json_write_sync(users, str(Path(f"tmp/{uuid}")))
+
+                await ctx.respond(
+                    file=File(str(Path(f"tmp/{uuid}")), filename="users.json")
                 )
-
-            json_write_sync(users, str(Path(f"tmp/{uuid}")))
-
-            await ctx.respond(file=File(str(Path(f"tmp/{uuid}")), filename="users.json"))
         else:
             logging.info(
                 f"User {guild_name(ctx.user)} tried to use /export but permission denied"
@@ -77,3 +83,70 @@ class Data(commands.Cog):
                     color=Color.fail,
                 ),
             )
+
+    @data.command(
+        name="migrate",
+        description="Мігрувати всіх користувачів до бази",
+        guild_ids=[config_get_sync("guild")],
+    )
+    @option(
+        "kind", description="Тип даних, які треба експортувати", choices=["Користувачі"]
+    )
+    async def data_migrate_cmd(self, ctx: ApplicationContext, kind: str):
+        await ctx.defer()
+        if ctx.user.id in await config_get("admins"):
+            logging.info(
+                f"Moderator {guild_name(ctx.user)} started migration of all members to the database"
+            )
+
+            if kind == "Користувачі":
+                for member in ctx.guild.members:
+                    if member.bot:
+                        continue
+                    if col_users.find_one({"user": member.id}) is None:
+                        user = {}
+                        defaults = await config_get("user", "defaults")
+
+                        user["user"] = member.id
+
+                        for key in defaults:
+                            user[key] = defaults[key]
+
+                        col_users.insert_one(document=user)
+                        logging.info(
+                            f"Added DB record for user {member.id} during migration"
+                        )
+
+            await ctx.respond(
+                embed=Embed(
+                    title="Міграцію завершено",
+                    description="Всім користувачам сервера було створено записи в базі даних.",
+                    color=Color.success,
+                )
+            )
+        else:
+            logging.info(
+                f"User {guild_name(ctx.user)} tried to use /migrate but permission denied"
+            )
+            await ctx.respond(
+                embed=Embed(
+                    title="Відмовлено в доступі",
+                    description="Здається, це команда лише для модераторів",
+                    color=Color.fail,
+                )
+            )
+            mod_role = ds_utils.get(
+                ctx.user.guild.roles, id=await config_get("moderator", "roles")
+            )
+            admin_chan = ds_utils.get(
+                ctx.user.guild.channels,
+                id=await config_get("adminchat", "channels", "text"),
+            )
+            await admin_chan.send(
+                content=f"{mod_role.mention}",
+                embed=Embed(
+                    title="Неавторизований запит",
+                    description=f"Користувач {ctx.user.mention} запитав у каналі {ctx.channel.mention} команду, до якої не повинен мати доступу/бачити.",
+                    color=Color.fail,
+                ),
+            )
diff --git a/main.py b/main.py
index 613c574..e58a2b1 100644
--- a/main.py
+++ b/main.py
@@ -42,7 +42,7 @@ def main():
     client.add_cog(Admin(client))
     client.add_cog(Analytics(client))
     client.add_cog(CustomChannels(client))
-    # client.add_cog(Data(client))
+    client.add_cog(Data(client))
     client.add_cog(Logger(client))
 
     try: