homelab/userbot/modules/misc/backup.py

import os
import shutil

from pyrogram import Client, enums, filters
from pyrogram.types import Message
from utils import config
from utils.db import db

# noinspection PyUnresolvedReferences
from utils.misc import modules_help, prefix
from utils.scripts import format_exc, restart

if config.db_type in ['mongodb', 'mongo']:
    import bson


def dump_mongo(collections, path, db_):
    for coll in collections:
        with open(os.path.join(path, f'{coll}.bson'), 'wb+') as f:
            for doc in db_[coll].find():
                f.write(bson.BSON.encode(doc))


def restore_mongo(path, db_):
    for coll in os.listdir(path):
        if coll.endswith('.bson'):
            with open(os.path.join(path, coll), 'rb+') as f:
                db_[coll.split('.')[0]].insert_many(bson.decode_all(f.read()))


@Client.on_message(filters.command(['backup', 'back'], prefix) & filters.me)
async def backup(client: Client, message: Message):
    """
    Backup the database
    """
    try:
        if not os.path.exists('backups/'):
            os.mkdir('backups/')

        await message.edit('<b>Backing up database...</b>', parse_mode=enums.ParseMode.HTML)

        if config.db_type in ['mongo', 'mongodb']:
            dump_mongo(db._database.list_collection_names(), 'backups/', db._database)
            return await message.edit(
                '<b>Database backed up to:</b> <code>backups/</code> folder',
                parse_mode=enums.ParseMode.HTML,
            )
        else:
            shutil.copy(config.db_name, f'backups/{config.db_name}')
            await client.send_document(
                'me',
                caption='<b>Database backup complete!\nType: </b>'
                '<code>.restore</code> in response to this message to restore the database.',
                document=f'backups/{config.db_name}',
                parse_mode=enums.ParseMode.HTML,
            )
            return await message.edit(
                '<b>Database backed up successfully! <code>(Check your favorites)</code></b>',
                parse_mode=enums.ParseMode.HTML,
            )
    except Exception as e:
        await message.edit(format_exc(e), parse_mode=enums.ParseMode.HTML)


@Client.on_message(filters.command(['restore', 'res'], prefix) & filters.me)
async def restore(_client: Client, message: Message):
    """
    Restore the database
    """
    try:
        await message.edit('<b>Restoring database...</b>', parse_mode=enums.ParseMode.HTML)
        if config.db_type in ['mongo', 'mongodb']:
            restore_mongo('backups/', db._database)
            return await message.edit(
                '<b>Database restored from:</b> <code>backups/</code> folder',
                parse_mode=enums.ParseMode.HTML,
            )
        else:
            if not message.reply_to_message or not message.reply_to_message.document:
                return await message.edit(
                    '<b>Reply to a document to restore the database.</b>',
                    parse_mode=enums.ParseMode.HTML,
                )
            elif not message.reply_to_message.document.file_name.casefold().endswith(('.db', '.sqlite', '.sqlite3')):
                return await message.edit(
                    '<b>Reply to a database file to restore the database.</b>',
                    parse_mode=enums.ParseMode.HTML,
                )
            await message.reply_to_message.download(f'backups/{message.reply_to_message.document.file_name}')
            shutil.copy(f'backups/{message.reply_to_message.document.file_name}', config.db_name)
            await message.edit(
                '<b>Database restored successfully!</b>',
                parse_mode=enums.ParseMode.HTML,
            )
            restart()
    except Exception as e:
        await message.edit(format_exc(e), parse_mode=enums.ParseMode.HTML)


@Client.on_message(filters.command(['backupmods', 'bms'], prefix) & filters.me)
async def backupmods(_client: Client, message: Message):
    """
    Backup the modules
    """
    try:
        if not os.path.exists('backups/'):
            os.mkdir('backups/')

        await message.edit('<b>Backing up modules...</b>', parse_mode=enums.ParseMode.HTML)

        from utils.misc import modules_help

        for mod in modules_help:
            if os.path.isfile(f'modules/custom_modules/{mod}.py'):
                with open(f'backups/{mod}.py', 'wb') as f, open(f'modules/custom_modules/{mod}.py', 'rb') as src:
                    f.write(src.read())
        await message.edit(
            text='<b>All modules backed up to:</b> <code>backups/</code> folder',
            parse_mode=enums.ParseMode.HTML,
        )
    except Exception as e:
        await message.edit(format_exc(e), parse_mode=enums.ParseMode.HTML)


@Client.on_message(filters.command(['backupmod', 'bm'], prefix) & filters.me)
async def backupmod(_client: Client, message: Message):
    """
    Backup the module
    """
    try:
        if not os.path.exists('backups/'):
            os.mkdir('backups/')

        try:
            mod = message.text.split(maxsplit=1)[1].split('.')[0]
        except Exception:
            return await message.edit(
                f'<b>Usage:</b> <code>{prefix}backupmod [module]</code>',
                parse_mode=enums.ParseMode.HTML,
            )

        await message.edit('<b>Backing up module...</b>', parse_mode=enums.ParseMode.HTML)

        if os.path.isfile(f'modules/custom_modules/{mod}.py'):
            with open(f'backups/{mod}.py', 'wb') as f, open(f'modules/custom_modules/{mod}.py', 'rb') as src:
                f.write(src.read())
        else:
            return await message.edit(
                f'<b>Module <code>{mod}</code> not found.</b>',
                parse_mode=enums.ParseMode.HTML,
            )
        await message.reply_document(
            document=f'backups/{mod}.py',
            caption=f'<b>Module <code>{mod}</code> backed up to:</b> <code>backups/</code> folder',
            parse_mode=enums.ParseMode.HTML,
        )
    except Exception as e:
        await message.edit(format_exc(e), parse_mode=enums.ParseMode.HTML)


@Client.on_message(filters.command(['restoremod', 'resmod'], prefix) & filters.me)
async def restoremod(_client: Client, message: Message):
    """
    Restore the module
    """
    try:
        if not os.path.exists('backups/'):
            os.mkdir('backups/')

        try:
            mod = message.text.split(maxsplit=1)[1].split('.')[0]
        except Exception:
            return await message.edit(
                f'<b>Usage:</b> <code>{prefix}restoremod [module]</code>',
                parse_mode=enums.ParseMode.HTML,
            )

        await message.edit('<b>Restoring module...</b>', parse_mode=enums.ParseMode.HTML)

        if os.path.isfile(f'backups/{mod}.py'):
            with open(f'modules/custom_modules/{mod}.py', 'wb') as f, open(f'backups/{mod}.py', 'rb') as src:
                f.write(src.read())
        else:
            return await message.edit(
                text='<b>Backup file <code>{mod}</code> not found</b>',
                parse_mode=enums.ParseMode.HTML,
            )
        await message.edit(
            f'<b>Module <code>{mod}</code> restored successfully!</b>',
            parse_mode=enums.ParseMode.HTML,
        )
        restart()
    except Exception as e:
        await message.edit(format_exc(e), parse_mode=enums.ParseMode.HTML)


@Client.on_message(filters.command(['restoremods', 'resmods'], prefix) & filters.me)
async def restoremods(_client: Client, message: Message):
    """
    Restore the modules
    """
    try:
        if not os.path.exists('backups/'):
            os.mkdir('backups/')

        await message.edit('<b>Restoring modules...</b>', parse_mode=enums.ParseMode.HTML)

        for mod in os.listdir('backups/'):
            if mod.endswith('.py'):
                if os.path.isfile(f'modules/{mod}'):
                    continue
                with open(f'modules/custom_modules/{mod}', 'wb') as f, open(f'backups/{mod}', 'rb') as src:
                    f.write(src.read())
        await message.edit(
            text='<b>All modules restored from:</b> <code>backups/</code> folder',
            parse_mode=enums.ParseMode.HTML,
        )
        restart()
    except Exception as e:
        await message.edit(format_exc(e), parse_mode=enums.ParseMode.HTML)


modules_help['backup'] = {
    'backup': '<b>Backup database</b>',
    'restore [reply]': '<b>Restore database</b>',
    'backupmod [name]': '<b>Backup mod</b>',
    'backupmods': '<b>Backup all mods</b>',
    'resmod [name]': '<b>Restore mod</b>',
    'resmods': '<b>Restore all mods</b>',
}