From 080fc78516c025195c708b1bd0de8b9d0810ae1c Mon Sep 17 00:00:00 2001 From: flifloo Date: Sat, 28 Mar 2020 14:04:20 +0100 Subject: [PATCH] Separation of backup and restore commands --- extensions/__init__.py | 1 + extensions/backup.py | 68 +----------------------------- extensions/restore.py | 96 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 98 insertions(+), 67 deletions(-) create mode 100644 extensions/restore.py diff --git a/extensions/__init__.py b/extensions/__init__.py index 7e544c5..3a6abab 100644 --- a/extensions/__init__.py +++ b/extensions/__init__.py @@ -2,3 +2,4 @@ from backup_bot import bot bot.load_extension("extensions.help") bot.load_extension("extensions.backup") +bot.load_extension("extensions.restore") diff --git a/extensions/backup.py b/extensions/backup.py index 805da3f..2b69cc8 100644 --- a/extensions/backup.py +++ b/extensions/backup.py @@ -1,10 +1,9 @@ from discord.ext import commands from backup_bot.logger import logger from os.path import isdir -from os import mkdir, remove +from os import mkdir import shelve from datetime import datetime -from requests import get from discord import File, Embed from collections import OrderedDict @@ -63,69 +62,6 @@ async def backup_cmd(ctx: commands.Context): await ctx.send(file=File(file_name + ".db", "backup.db")) -@commands.command("restore") -async def restore_cmd(ctx: commands.Context): - if len(ctx.message.attachments) != 1: - await ctx.send("No backup file given ! \N{cross mark}") - else: - embed = Embed(title="Restore", description="In progress... \N{hourglass}") - msg = await ctx.send(embed=embed) - - file = get(ctx.message.attachments[0].url, stream=True) - file_name = f"backup/{ctx.message.author.id}" - with open(file_name + ".db", "w+b") as f: - for i in file.iter_content(): - f.write(i) - with shelve.open(file_name) as file: - categories = {} - for c in file["categories"]: - categories[c] = await ctx.guild.create_category(name=file["categories"][c]["name"], - reason=f"Backup restore by {ctx.message.author}") - for c in file["channels"]: - embed_field_name = file["channels"][c]["name"] - category = None - if file["channels"][c]["category_id"]: - category = categories[file["channels"][c]["category_id"]] - embed_field_name = f"{category.name} > {embed_field_name}" - - embed = msg.embeds[0] - if len(embed.fields) != 0: - embed.set_field_at(-1, name=embed.fields[-1].name, value="\N{check mark}", inline=False) - embed.add_field(name=embed_field_name, value="\N{hourglass}", inline=False) - await msg.edit(embed=embed) - - chan = await ctx.guild.create_text_channel(name=file["channels"][c]["name"], - category=category, - topic=file["channels"][c]["topic"], - slowmode_delay=file["channels"][c]["slowmode_delay"], - nsfw=file["channels"][c]["nsfw"], - reason=f"Backup restore by {ctx.message.author}") - hook = await chan.create_webhook(name="BackupBot", - avatar=None, - reason=f"Backup restore by {ctx.message.author}") - for m in file["channels"][c]["messages"][::-1]: - user = file["users"][m["author_id"]] - edit = "" - if m["edited_at"]: - edit = f", edited at: {m['edited_at']}" - content = f"`created: {m['created_at']}{edit}`" + "\n" + m["content"] - avatar = None - if user["avatar"]: - avatar = f"https://cdn.discordapp.com/avatars/{m['author_id']}/{user['avatar']}.webp" - await hook.send(content=content, - username=f"{user['display_name']} ({user['name']}#{user['discriminator']})", - avatar_url=avatar, - files=None, - embeds=m["embeds"]) - await hook.delete() - - remove(file_name + ".db") - embed = msg.embeds[0] - embed.set_field_at(-1, name=embed.fields[-1].name, value="\N{check mark}", inline=False) - embed.description = "Finish ! \N{check mark}" - await msg.edit(embed=embed) - - def setup(bot: commands.Bot): logger.info(f"Loading of {extension_name} extension") if not isdir("backup"): @@ -133,7 +69,6 @@ def setup(bot: commands.Bot): mkdir("backup") try: bot.add_command(backup_cmd) - bot.add_command(restore_cmd) except Exception as e: logger.error(f"Error loading extension {extension_name}: {e}") else: @@ -144,7 +79,6 @@ def teardown(bot: commands.Bot): logger.info(f"Unloading of {extension_name} extension") try: bot.remove_command("backup") - bot.remove_command("restore") except Exception as e: logger.error(f"Error unloading extension {extension_name}: {e}") else: diff --git a/extensions/restore.py b/extensions/restore.py new file mode 100644 index 0000000..643f3ca --- /dev/null +++ b/extensions/restore.py @@ -0,0 +1,96 @@ +from discord.ext import commands +from backup_bot.logger import logger +from os.path import isdir +from os import mkdir, remove +import shelve +from requests import get +from discord import Embed + +extension_name = "restore" +logger = logger.getChild(extension_name) + + +@commands.command("restore") +async def restore_cmd(ctx: commands.Context): + if len(ctx.message.attachments) != 1: + await ctx.send("No backup file given ! \N{cross mark}") + else: + embed = Embed(title="Restore", description="In progress... \N{hourglass}") + msg = await ctx.send(embed=embed) + + file = get(ctx.message.attachments[0].url, stream=True) + file_name = f"backup/{ctx.message.author.id}" + with open(file_name + ".db", "w+b") as f: + for i in file.iter_content(): + f.write(i) + with shelve.open(file_name) as file: + categories = {} + for c in file["categories"]: + categories[c] = await ctx.guild.create_category(name=file["categories"][c]["name"], + reason=f"Backup restore by {ctx.message.author}") + for c in file["channels"]: + embed_field_name = file["channels"][c]["name"] + category = None + if file["channels"][c]["category_id"]: + category = categories[file["channels"][c]["category_id"]] + embed_field_name = f"{category.name} > {embed_field_name}" + + embed = msg.embeds[0] + if len(embed.fields) != 0: + embed.set_field_at(-1, name=embed.fields[-1].name, value="\N{check mark}", inline=False) + embed.add_field(name=embed_field_name, value="\N{hourglass}", inline=False) + await msg.edit(embed=embed) + + chan = await ctx.guild.create_text_channel(name=file["channels"][c]["name"], + category=category, + topic=file["channels"][c]["topic"], + slowmode_delay=file["channels"][c]["slowmode_delay"], + nsfw=file["channels"][c]["nsfw"], + reason=f"Backup restore by {ctx.message.author}") + hook = await chan.create_webhook(name="BackupBot", + avatar=None, + reason=f"Backup restore by {ctx.message.author}") + for m in file["channels"][c]["messages"][::-1]: + user = file["users"][m["author_id"]] + edit = "" + if m["edited_at"]: + edit = f", edited at: {m['edited_at']}" + content = f"`created: {m['created_at']}{edit}`" + "\n" + m["content"] + avatar = None + if user["avatar"]: + avatar = f"https://cdn.discordapp.com/avatars/{m['author_id']}/{user['avatar']}.webp" + await hook.send(content=content, + username=f"{user['display_name']} ({user['name']}#{user['discriminator']})", + avatar_url=avatar, + files=None, + embeds=m["embeds"]) + await hook.delete() + + remove(file_name + ".db") + embed = msg.embeds[0] + embed.set_field_at(-1, name=embed.fields[-1].name, value="\N{check mark}", inline=False) + embed.description = "Finish ! \N{check mark}" + await msg.edit(embed=embed) + + +def setup(bot: commands.Bot): + logger.info(f"Loading of {extension_name} extension") + if not isdir("backup"): + logger.info(f"Create backup folder") + mkdir("backup") + try: + bot.add_command(restore_cmd) + except Exception as e: + logger.error(f"Error loading extension {extension_name}: {e}") + else: + logger.info(f"Extension {extension_name} load successful") + + +def teardown(bot: commands.Bot): + logger.info(f"Unloading of {extension_name} extension") + try: + bot.remove_command("restore") + except Exception as e: + logger.error(f"Error unloading extension {extension_name}: {e}") + else: + logger.info(f"Extension {extension_name} unload successful")