Discord/cogs/data.py

import logging
from os import makedirs
from pathlib import Path
from uuid import uuid4

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 libbot import config_get
from libbot.pycord.classes import PycordBot
from libbot.sync import config_get as sync_config_get
from libbot.sync import json_write as sync_json_write

from classes.holo_user import HoloUser
from enums.colors import Color
from modules.database import col_users
from modules.utils_sync import guild_name

logger = logging.getLogger(__name__)


class Data(commands.Cog):
    def __init__(self, client: PycordBot):
        self.client = client

    data = SlashCommandGroup("data", "Керування даними користувачів")

    @data.command(
        name="export",
        description="Експортувати дані",
        guild_ids=[sync_config_get("guild")],
    )
    @option(
        "kind", description="Тип даних, які треба експортувати", choices=["Користувачі"]
    )
    async def data_export_cmd(self, ctx: ApplicationContext, kind: str):
        await ctx.defer()
        holo_user = HoloUser(ctx.author)

        # Return if the user is not an owner and not in the council
        if (ctx.user.id not in self.client.owner_ids) and not (
            await HoloUser.is_council(ctx.author)
        ):
            logging.info(
                "User %s tried to use /export but permission denied",
                guild_name(ctx.user),
            )

            await ctx.respond(
                embed=Embed(
                    title="Відмовлено в доступі",
                    description="Здається, це команда лише для модераторів",
                    color=Color.fail,
                )
            )

            mod_role = ds_utils.get(
                ctx.user.guild.roles, id=await config_get("moderators", "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,
                ),
            )

            return

        logging.info("Moderator %s exported current users list", guild_name(ctx.user))

        makedirs("tmp", exist_ok=True)

        uuid = str(uuid4())

        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,
                    }
                )

            sync_json_write(users, Path(f"tmp/{uuid}"))

            await ctx.respond(file=File(Path(f"tmp/{uuid}"), filename="users.json"))

    @data.command(
        name="migrate",
        description="Мігрувати всіх користувачів до бази",
        guild_ids=[sync_config_get("guild")],
    )
    @option(
        "kind", description="Тип даних, які треба експортувати", choices=["Користувачі"]
    )
    async def data_migrate_cmd(self, ctx: ApplicationContext, kind: str):
        await ctx.defer()

        # Return if the user is not an owner and not in the council
        if (ctx.user.id not in self.client.owner_ids) and not (
            await HoloUser.is_council(ctx.author)
        ):
            logging.info(
                "User %s tried to use /migrate but permission denied",
                guild_name(ctx.user),
            )

            await ctx.respond(
                embed=Embed(
                    title="Відмовлено в доступі",
                    description="Здається, це команда лише для модераторів",
                    color=Color.fail,
                )
            )

            mod_role = ds_utils.get(
                ctx.user.guild.roles, id=await config_get("moderators", "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,
                ),
            )

            return

        logging.info(
            "Moderator %s started migration of all members to the database",
            guild_name(ctx.user),
        )

        if kind == "Користувачі":
            for member in ctx.guild.members:
                if member.bot:
                    continue

                if (await 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]

                    await col_users.insert_one(document=user)

                    logging.info(
                        "Added DB record for user %s during migration", member.id
                    )

        await ctx.respond(
            embed=Embed(
                title="Міграцію завершено",
                description="Всім користувачам сервера було створено записи в базі даних.",
                color=Color.success,
            )
        )


def setup(client: PycordBot):
    client.add_cog(Data(client))