Discord/cogs/data.py

import logging
from os import makedirs
from pathlib import Path
from typing import Union, List, Dict, Any
from uuid import uuid4

from discord import ApplicationContext, Embed, File, option, Role, TextChannel
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 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: PycordBot = client

    data: SlashCommandGroup = 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) -> None:
        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 /export but permission denied",
                guild_name(ctx.user),
            )

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

            mod_role: Union[Role, None] = ds_utils.get(
                ctx.user.guild.roles, id=await config_get("moderators", "roles")
            )
            admin_chan: Union[TextChannel, None] = ds_utils.get(
                ctx.user.guild.channels,
                id=await config_get("adminchat", "channels", "text"),
            )

            await admin_chan.send(
                content="" if mod_role is None else 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 = str(uuid4())

        if kind == "Користувачі":
            users: List[Dict[str, Any]] = []

            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) -> None:
        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: Union[Role, None] = ds_utils.get(
                ctx.user.guild.roles, id=await config_get("moderators", "roles")
            )
            admin_chan: Union[TextChannel, None] = ds_utils.get(
                ctx.user.guild.channels,
                id=await config_get("adminchat", "channels", "text"),
            )

            if admin_chan is not None:
                await admin_chan.send(
                    content="" if mod_role is None else 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: Dict[str, Any] = {}
                    defaults: Dict[str, Any] = 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) -> None:
    client.add_cog(Data(client))