TelegramPoster/plugins/commands/photos.py

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