SeaCogs/backup/backup.py

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

import contextlib
import json
import re

from red_commons.logging import getLogger
from redbot.cogs.downloader import errors
from redbot.cogs.downloader.converters import InstalledCog
from redbot.core import commands
from redbot.core.bot import Red
from redbot.core.utils.chat_formatting import error, humanize_list, text_to_file


# pylint: disable=protected-access
class Backup(commands.Cog):
    """A utility to make reinstalling repositories and cogs after migrating the bot far easier."""

    __author__ = ["SeaswimmerTheFsh"]
    __version__ = "1.1.0"
    __documentation__ = "https://seacogs.coastalcommits.com/backup/"

    def __init__(self, bot: Red):
        super().__init__()
        self.bot = bot
        self.logger = getLogger("red.SeaCogs.Backup")

    def format_help_for_context(self, ctx: commands.Context) -> str:
        pre_processed = super().format_help_for_context(ctx) or ""
        n = "\n" if "\n\n" not in pre_processed else ""
        text = [
            f"{pre_processed}{n}",
            f"Cog Version: **{self.__version__}**",
            f"Author: {humanize_list(self.__author__)}",
            f"Documentation: {self.__documentation__}",
        ]
        return "\n".join(text)

    @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)

        export_data = []

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

            cogs = await downloader.installed_cogs()

            for cog in cogs:
                if cog.repo_name == repo.name:
                    cog_dict = {
                        "name": cog.name,
                        # "loaded": cog.name in ctx.bot.extensions.keys(),
                        # this functionality was planned but never implemented due to Red limitations
                        # and the possibility of restoration functionality being added to Core
                        "pinned": cog.pinned,
                        "commit": cog.commit,
                    }
                    repo_dict["cogs"].append(cog_dict)

            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):
        """Import your installed repositories and cogs from an export file."""
        export = None
        if ctx.message.attachments:
            try:
                export = json.loads(await ctx.message.attachments[0].read())
            except json.JSONDecodeError:
                await ctx.send(error("Invalid JSON in message attachments."))
        elif ctx.message.reference and hasattr(ctx.message.reference, 'resolved'):
            if ctx.message.reference.resolved.attachments:
                try:
                    export = json.loads(await ctx.message.reference.resolved.attachments[0].read())
                except json.JSONDecodeError:
                    await ctx.send(error("Invalid JSON in referenced message attachments."))
        if export is None:
            await ctx.send(error("Please provide a valid JSON export file."))
            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

        all_repos = list(downloader._repo_manager.repos)

        for repo in export:
            if repo["name"] not in [r.name for r in all_repos]:
                try:
                    await downloader._repo_manager.add_repo(
                        repo["url"], repo["name"], repo["branch"]
                    )
                except errors.ExistingGitRepo:
                    pass

            for cog in repo["cogs"]:
                try:
                    await downloader._cog_install_interface.install_cog(
                        cog["name"], cog["commit"]
                    )
                except errors.CogNotFoundError:
                    pass
                except errors.DownloaderError:
                    pass