Kuby/commandes/backup_restore.py

# commandes/backup_restore.py
import os
import json
import zipfile
import tempfile
import shutil
from typing import List
import disnake
from disnake.ext import commands

BACKUPS_ROOT = "data/security_backups"

def ensure_dir(path: str):
    os.makedirs(path, exist_ok=True)

class BackupRestore(commands.Cog):
    """Restauration granulaire d'une sauvegarde."""

    def __init__(self, bot: commands.Bot):
        self.bot = bot

    async def check_whitelisted(self, guild_id: int, user_id: int) -> bool:
        """Vérifie si un utilisateur est whitelisté."""
        # Recherche dynamique d’un cog ayant une méthode is_whitelisted
        for cog in self.bot.cogs.values():
            if hasattr(cog, "is_whitelisted"):
                return cog.is_whitelisted(guild_id, user_id)
        raise RuntimeError("Cog Whitelist non chargé ou mal nommé !")

    async def list_backup_elements(self, backup_path: str) -> List[str]:
        """Liste les éléments exportables dans une archive ou dossier de backup."""
        elements = []
        if os.path.isdir(backup_path):
            for f in os.listdir(backup_path):
                if f.endswith(".json") or os.path.isdir(os.path.join(backup_path, f)):
                    elements.append(f.replace(".json", ""))
        elif os.path.isfile(backup_path) and backup_path.endswith(".zip"):
            with zipfile.ZipFile(backup_path, "r") as zf:
                for f in zf.namelist():
                    base = f.split("/")[0]
                    if base and base not in elements:
                        elements.append(base)
        return elements

    async def element_autocomplete(self, interaction: disnake.ApplicationCommandInteraction, current: str):
        """Complétion automatique pour l'option 'element'."""
        guild = interaction.guild
        if guild is None:
            return []

        ensure_dir(BACKUPS_ROOT)
        files = [f for f in os.listdir(BACKUPS_ROOT) if f.startswith(str(guild.id))]
        files.sort(reverse=True)

        if not files:
            return []

        last_backup = files[0]
        backup_path = os.path.join(BACKUPS_ROOT, last_backup)
        elements = await self.list_backup_elements(backup_path)

        return [e for e in elements if current.lower() in e.lower()][:25]

    async def restore_roles(self, guild: disnake.Guild, backup_path: str):
        """Restaure les rôles à partir du backup."""
        roles_file = os.path.join(backup_path, "roles.json")
        if os.path.isfile(roles_file):
            with open(roles_file, "r", encoding="utf-8") as f:
                roles_data = json.load(f)
            for r in roles_data:
                try:
                    await guild.create_role(
                        name=r["name"],
                        permissions=disnake.Permissions(r["permissions"]),
                        color=disnake.Color(r.get("color", 0)),
                        hoist=r.get("hoist", False),
                        mentionable=r.get("mentionable", False),
                        reason="Restauration backup"
                    )
                except disnake.Forbidden:
                    continue

    async def restore_channels(self, guild: disnake.Guild, backup_path: str):
        """Restaure les channels à partir du backup."""
        channels_file = os.path.join(backup_path, "channels.json")
        if os.path.isfile(channels_file):
            with open(channels_file, "r", encoding="utf-8") as f:
                channels_data = json.load(f)
            for c in channels_data:
                try:
                    if c["type"] == 0:  # text
                        await guild.create_text_channel(
                            name=c["name"],
                            category=guild.get_channel(c.get("category_id")),
                            topic=c.get("topic", ""),
                            nsfw=c.get("nsfw", False),
                            reason="Restauration backup"
                        )
                    elif c["type"] == 2:  # voice
                        await guild.create_voice_channel(
                            name=c["name"],
                            category=guild.get_channel(c.get("category_id")),
                            bitrate=c.get("bitrate", 64000),
                            user_limit=c.get("user_limit", 0),
                            reason="Restauration backup"
                        )
                except disnake.Forbidden:
                    continue

    async def restore_messages(self, channel: disnake.TextChannel, backup_path: str):
        """Restaure les messages à partir du backup (solution ZIP incluse)."""
        if backup_path.endswith(".zip"):
            temp_dir = tempfile.mkdtemp()
            try:
                with zipfile.ZipFile(backup_path, "r") as zip_ref:
                    zip_ref.extractall(temp_dir)
                messages_file = os.path.join(temp_dir, "messages.json")
                if not os.path.exists(messages_file):
                    return
                with open(messages_file, "r", encoding="utf-8") as f:
                    messages = json.load(f)
                for msg in messages:
                    await channel.send(msg["content"])
            finally:
                shutil.rmtree(temp_dir)
        else:
            messages_file = os.path.join(backup_path, "messages.json")
            if not os.path.exists(messages_file):
                return
            with open(messages_file, "r", encoding="utf-8") as f:
                messages = json.load(f)
            for msg in messages:
                await channel.send(msg["content"])

    @commands.slash_command(name="granulaire", description="Restaure un élément précis d'une sauvegarde")
    async def restore_granulaire(
        self,
        interaction: disnake.ApplicationCommandInteraction,
        backup: str = commands.Param(description="Nom ou ID de la sauvegarde"),
        element: str = commands.Param(description="Élément à restaurer (roles, channels, messages...)")
    ):
        guild = interaction.guild
        if guild is None:
            await interaction.response.send_message("Cette commande doit être utilisée dans un serveur.", ephemeral=True)
            return

        if not await self.check_whitelisted(guild.id, interaction.user.id):
            await interaction.response.send_message(
                "❌ Vous n'êtes pas autorisé à utiliser cette commande.",
                ephemeral=True
            )
            return

        backup_path = os.path.join(BACKUPS_ROOT, backup)
        if not os.path.exists(backup_path):
            await interaction.response.send_message(f"❌ La sauvegarde `{backup}` est introuvable.", ephemeral=True)
            return

        await interaction.response.defer(ephemeral=True)

        available_elements = await self.list_backup_elements(backup_path)
        if element not in available_elements:
            await interaction.followup.send(
                f"❌ Élément `{element}` introuvable dans la sauvegarde.",
                ephemeral=True
            )
            return

        # Détection si c’est un zip pour extraire temporairement
        temp_path = backup_path
        if backup_path.endswith(".zip"):
            temp_path = tempfile.mkdtemp()
            with zipfile.ZipFile(backup_path, "r") as zip_ref:
                zip_ref.extractall(temp_path)

        try:
            if element == "roles":
                await self.restore_roles(guild, temp_path)
            elif element == "channels":
                await self.restore_channels(guild, temp_path)
            elif element == "messages":
                # Restaure messages dans tous les channels textuels ?
                for ch in guild.text_channels:
                    await self.restore_messages(ch, temp_path)
            else:
                await interaction.followup.send(f"⚠️ Élément `{element}` non géré par la restauration automatique.", ephemeral=True)
                return
        finally:
            if backup_path.endswith(".zip"):
                shutil.rmtree(temp_path)

        await interaction.followup.send(
            f"✅ Restauration granulaire de `{element}` depuis la sauvegarde `{backup}` terminée.",
            ephemeral=True
        )

    @restore_granulaire.autocomplete("element")
    async def element_autocomplete_decorator(self, interaction: disnake.ApplicationCommandInteraction, current: str):
        return await self.element_autocomplete(interaction, current)

async def setup(bot: commands.Bot):
    await bot.add_cog(BackupRestore(bot))