SeaCogs/backup/backup.py

#   _____                         _
#  / ____|                       (_)
# | (___   ___  __ _ _____      ___ _ __ ___  _ __ ___   ___ _ __
#  \___ \ / _ \/ _` / __\ \ /\ / / | '_ ` _ \| '_ ` _ \ / _ \ '__|
#  ____) |  __/ (_| \__ \\ V  V /| | | | | | | | | | | |  __/ |
# |_____/ \___|\__,_|___/ \_/\_/ |_|_| |_| |_|_| |_| |_|\___|_|

import json
import re

from redbot.core import commands
from redbot.core.bot import Red
from redbot.core.utils.chat_formatting import error, text_to_file

class Backup(commands.Cog):
    """A utility to make reinstalling repositories and cogs after migrating the bot far easier."""

    __author__ = "SeaswimmerTheFsh"
    __version__ = "1.0.0"

    def __init__(self, bot: Red):
        super().__init__()
        self.bot = bot

    @commands.group(autohelp=True)
    @commands.is_owner()
    async def backup(self, ctx: commands.Context):
        """Backup your installed cogs."""

    @backup.command(name='export')
    @commands.is_owner()
    async def backup_export(self, ctx: commands.Context):
        """Export your installed repositories and cogs to a file."""
        downloader = ctx.bot.get_cog("Downloader")
        if downloader is None:
            await ctx.send(error(f"You do not have the `Downloader` cog loaded. Please run `{ctx.prefix}load downloader` and try again."))
            return

        all_repos = list(downloader._repo_manager.repos) # pylint: disable=protected-access

        export_data = []

        for repo in all_repos:
            repo_dict = {
                "name": repo.name,
                "url": repo.url,
                "cogs": []
            }

            cogs = await downloader.installed_cogs()

            for cog in cogs:
                if cog.repo_name == repo.name:
                    repo_dict["cogs"].append(cog.name)

            export_data.append(repo_dict)

        await ctx.send(file=text_to_file(json.dumps(export_data, indent=4), 'backup.json'))

    @backup.command(name='import')
    @commands.is_owner()
    async def backup_import(self, ctx: commands.Context, *, export: str = None):
        """Import your installed repositories and cogs from an export."""
        if export is None and not len(ctx.message.attachments) == 0:
            export = await ctx.message.attachments[0].read()
        try:
            export = json.loads(self.extract_json_from_codeblock(export))
        except json.JSONDecodeError:
            await ctx.send(error("Please provide a valid JSON export."))
            return

        downloader = ctx.bot.get_cog("Downloader")
        if downloader is None:
            await ctx.send(error(f"You do not have the `Downloader` cog loaded. Please run `{ctx.prefix}load downloader` and try again."))
            return

    def extract_json_from_codeblock(text):
        match = re.match(r"```(?:\w+)?\n(.*?)\n```", text, re.DOTALL)
        if match:
            return match.group(1).strip()
        else:
            return text