import asyncio from glob import iglob from os import getcwd, makedirs, path, remove from shutil import disk_usage, rmtree from traceback import format_exc from uuid import uuid4 from zipfile import ZipFile from convopyro import listen_message from pyrogram import filters from pyrogram.types import Message from classes.poster_client import PosterClient from modules.api_client import upload_pic from modules.app import app, users_with_context from modules.logger import logWrite from modules.utils import configGet, extract_and_save, locale @app.on_message(~filters.scheduled & filters.command(["import"], prefixes=["", "/"])) async def cmd_import(app: PosterClient, msg: Message): if msg.from_user.id in app.admins: global users_with_context if msg.from_user.id not in users_with_context: users_with_context.append(msg.from_user.id) else: return await msg.reply_text( locale("import_request", "message", locale=msg.from_user.language_code) ) answer = await listen_message(app, msg.chat.id, timeout=600) users_with_context.remove(msg.from_user.id) if answer is None: await msg.reply_text( locale("import_ignored", "message", locale=msg.from_user.language_code), quote=True, ) return if answer.text == "/cancel": await answer.reply_text( locale("import_abort", "message", locale=msg.from_user.language_code) ) return if answer.document is None: await answer.reply_text( locale( "import_invalid_media", "message", locale=msg.from_user.language_code, ), quote=True, ) return if answer.document.mime_type != "application/zip": await answer.reply_text( locale( "import_invalid_mime", "message", locale=msg.from_user.language_code ), quote=True, ) return if disk_usage(getcwd())[2] < (answer.document.file_size) * 3: await msg.reply_text( locale( "import_too_big", "message", locale=msg.from_user.language_code ).format( answer.document.file_size // (2**30), disk_usage(getcwd())[2] // (2**30), ) ) return tmp_dir = str(uuid4()) logWrite( f"Importing '{answer.document.file_name}' file {answer.document.file_size} bytes big (TMP ID {tmp_dir})" ) makedirs(path.join(configGet("tmp", "locations"), tmp_dir), exist_ok=True) tmp_path = path.join(configGet("tmp", "locations"), answer.document.file_id) downloading = await answer.reply_text( locale("import_downloading", "message", locale=msg.from_user.language_code), quote=True, ) await app.download_media(answer, file_name=tmp_path) await downloading.edit( locale("import_unpacking", "message", locale=msg.from_user.language_code) ) try: with ZipFile(tmp_path, "r") as handle: tasks = [ extract_and_save( handle, name, path.join(configGet("tmp", "locations"), tmp_dir) ) for name in handle.namelist() ] _ = await asyncio.gather(*tasks) except Exception as exp: logWrite( f"Could not import '{answer.document.file_name}' due to {exp}: {format_exc}" ) await answer.reply_text( locale( "import_unpack_error", "message", locale=msg.from_user.language_code ).format(exp, format_exc()) ) return logWrite(f"Downloaded '{answer.document.file_name}' - awaiting upload") await downloading.edit( locale("import_uploading", "message", locale=msg.from_user.language_code) ) remove(tmp_path) for filename in iglob( path.join(configGet("tmp", "locations"), tmp_dir) + "**/**", recursive=True ): if not path.isfile(filename): continue # upload filename uploaded = await upload_pic(filename) if uploaded[0] is False: logWrite( f"Could not upload '{filename}' from '{path.join(configGet('tmp', 'locations'), tmp_dir)}'. Duplicates: {str(uploaded[1])}", debug=True, ) if len(uploaded[1]) > 0: await msg.reply_text( locale( "import_upload_error_duplicate", "message", locale=msg.from_user.language_code, ).format(path.basename(filename)), disable_notification=True, ) else: await msg.reply_text( locale( "import_upload_error_other", "message", locale=msg.from_user.language_code, ).format(path.basename(filename)), disable_notification=True, ) else: logWrite( f"Uploaded '{filename}' from '{path.join(configGet('tmp', 'locations'), tmp_dir)}' and got ID {uploaded[2]}", debug=True, ) await downloading.delete() logWrite( f"Removing '{path.join(configGet('tmp', 'locations'), tmp_dir)}' after uploading", debug=True, ) rmtree(path.join(configGet("tmp", "locations"), tmp_dir), ignore_errors=True) await answer.reply_text( locale("import_finished", "message", locale=msg.from_user.language_code), quote=True, ) return @app.on_message(~filters.scheduled & filters.command(["export"], prefixes=["", "/"])) async def cmd_export(app: PosterClient, msg: Message): if msg.from_user.id in app.admins: pass @app.on_message(~filters.scheduled & filters.command(["remove"], prefixes=["", "/"])) async def cmd_remove(app: PosterClient, msg: Message): if msg.from_user.id in app.admins: pass @app.on_message(~filters.scheduled & filters.command(["purge"], prefixes=["", "/"])) async def cmd_purge(app: PosterClient, msg: Message): if msg.from_user.id in app.admins: pass