From 80c012438edafbfe4e5ecc1eddde69656e02be2a Mon Sep 17 00:00:00 2001 From: flifloo Date: Thu, 30 Jul 2020 11:56:31 +0200 Subject: [PATCH] Integrate backup & restore extensions --- extensions/__init__.py | 2 + extensions/backup.py | 130 ++++++++++++++++++++----------------- extensions/restore.py | 144 ++++++++++++++++++++++------------------- requirements.txt | 3 + 4 files changed, 154 insertions(+), 125 deletions(-) diff --git a/extensions/__init__.py b/extensions/__init__.py index 0b1a35f..0f0070c 100644 --- a/extensions/__init__.py +++ b/extensions/__init__.py @@ -8,3 +8,5 @@ bot.load_extension("extensions.reminder") bot.load_extension("extensions.greetings") bot.load_extension("extensions.presentation") bot.load_extension("extensions.rorec") +bot.load_extension("extensions.backup") +bot.load_extension("extensions.restore") diff --git a/extensions/backup.py b/extensions/backup.py index 111f53e..42e9317 100644 --- a/extensions/backup.py +++ b/extensions/backup.py @@ -1,84 +1,96 @@ -from discord.ext import commands -from backup_bot.logger import logger -from os.path import isdir -from os import mkdir import shelve -from datetime import datetime -from discord import File, Embed from collections import OrderedDict +from datetime import datetime +from os import mkdir +from os.path import isdir + +from discord import File, Embed +from discord.ext import commands + +from administrator.logger import logger + extension_name = "backup" logger = logger.getChild(extension_name) -@commands.command("backup") -async def backup_cmd(ctx: commands.Context): - embed = Embed(title="Backup", description="In progress... \N{hourglass}") - msg = await ctx.send(embed=embed) - file_name = f"backup/{datetime.now().strftime('%d-%m-%Y %H:%M')}" - with shelve.open(file_name, writeback=True) as file: - file["channels"] = OrderedDict() - file["users"] = OrderedDict() - file["categories"] = OrderedDict() - for c in ctx.guild.text_channels: - embed_field_name = c.name - if c.category: - embed_field_name = f"{c.category} > {embed_field_name}" - if c.category_id not in file["categories"]: - file["categories"][c.category_id] = {"name": c.category.name, - "position": c.category.position, - "nsfw": c.category.is_nsfw()} - 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) - file["channels"][c.id] = {"name": c.name, - "id": c.id, - "category_id": c.category_id, - "topic": c.topic, - "position": c.position, - "slowmode_delay": c.slowmode_delay, - "nsfw": c.is_nsfw(), - "messages": []} - async for m in c.history(limit=None): - if m.author.id not in file["users"]: - file["users"][m.author.id] = {"name": m.author.name, - "discriminator": m.author.discriminator, - "display_name": m.author.display_name, - "avatar": m.author.avatar} - file["channels"][c.id]["messages"].append({"author_id": m.author.id, - "content": m.content, - "embeds": m.embeds, - # "attachments": m.attachments, - "pinned": m.pinned, - "reactions": m.reactions, - "created_at": m.created_at, - "edited_at": m.edited_at}) - 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) - await ctx.send(file=File(file_name + ".db", "backup.db")) +class Backup(commands.Cog): + def __init__(self, bot: commands.Bot): + self.bot = bot + + def description(self): + return "Backup all message on the guild" + + @commands.group("backup", pass_context=True) + @commands.guild_only() + @commands.has_permissions(manage_messages=True) + async def backup(self, ctx: commands.Context): + embed = Embed(title="Backup", description="In progress... \N{hourglass}") + msg = await ctx.send(embed=embed) + file_name = f"backup/{datetime.now().strftime('%d-%m-%Y %H:%M')}" + with shelve.open(file_name, writeback=True) as file: + file["channels"] = OrderedDict() + file["users"] = OrderedDict() + file["categories"] = OrderedDict() + for c in ctx.guild.text_channels: + embed_field_name = c.name + if c.category: + embed_field_name = f"{c.category} > {embed_field_name}" + if c.category_id not in file["categories"]: + file["categories"][c.category_id] = {"name": c.category.name, + "position": c.category.position, + "nsfw": c.category.is_nsfw()} + 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) + file["channels"][c.id] = {"name": c.name, + "id": c.id, + "category_id": c.category_id, + "topic": c.topic, + "position": c.position, + "slowmode_delay": c.slowmode_delay, + "nsfw": c.is_nsfw(), + "messages": []} + async for m in c.history(limit=None): + if m.author.id not in file["users"]: + file["users"][m.author.id] = {"name": m.author.name, + "discriminator": m.author.discriminator, + "display_name": m.author.display_name, + "avatar": m.author.avatar} + file["channels"][c.id]["messages"].append({"author_id": m.author.id, + "content": m.content, + "embeds": m.embeds, + # "attachments": m.attachments, + "pinned": m.pinned, + "reactions": m.reactions, + "created_at": m.created_at, + "edited_at": m.edited_at}) + 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) + await ctx.send(file=File(file_name + ".db", "backup.db")) -def setup(bot: commands.Bot): +def setup(bot): logger.info(f"Loading...") if not isdir("backup"): logger.info(f"Create backup folder") mkdir("backup") try: - bot.add_command(backup_cmd) + bot.add_cog(Backup(bot)) except Exception as e: logger.error(f"Error loading: {e}") else: logger.info(f"Load successful") -def teardown(bot: commands.Bot): +def teardown(bot): logger.info(f"Unloading...") try: - bot.remove_command("backup") + bot.remove_cog("Backup") except Exception as e: logger.error(f"Error unloading: {e}") else: diff --git a/extensions/restore.py b/extensions/restore.py index c1c0d4b..75a268b 100644 --- a/extensions/restore.py +++ b/extensions/restore.py @@ -1,95 +1,107 @@ -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 os import mkdir, remove +from os.path import isdir + from discord import Embed +from discord.ext import commands +from requests import get + +from administrator.logger import logger + 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) +class Restore(commands.Cog): + def __init__(self, bot: commands.Bot): + self.bot = bot - 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}" + def description(self): + return "Restore a backup of guild's messages" - 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) + @commands.group("restore", pass_context=True) + @commands.guild_only() + @commands.has_permissions(manage_messages=True) + async def restore(self, 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) - 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() + 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}" - 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) + 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): +def setup(bot): logger.info(f"Loading...") if not isdir("backup"): logger.info(f"Create backup folder") mkdir("backup") try: - bot.add_command(restore_cmd) + bot.add_cog(Restore(bot)) except Exception as e: logger.error(f"Error loading: {e}") else: logger.info(f"Load successful") -def teardown(bot: commands.Bot): +def teardown(bot): logger.info(f"Unloading...") try: - bot.remove_command("restore") + bot.remove_cog("Restore") except Exception as e: logger.error(f"Error unloading: {e}") else: diff --git a/requirements.txt b/requirements.txt index 3bc624a..b64b776 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,14 @@ aiohttp==3.6.2 async-timeout==3.0.1 attrs==19.3.0 +certifi==2020.6.20 chardet==3.0.4 discord==1.0.1 discord.py==1.3.4 idna==2.10 multidict==4.7.6 +requests==2.24.0 SQLAlchemy==1.3.18 +urllib3==1.25.10 websockets==8.1 yarl==1.4.2