TelegramPoster/plugins/handlers/submission.py

from datetime import datetime
from os import makedirs, path, sep
from traceback import format_exc
from uuid import uuid4

from pyrogram import filters
from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup, Message
from pyrogram.enums.chat_action import ChatAction
from classes.exceptions import SubmissionDuplicatesError
from classes.poster_client import PosterClient
from classes.user import PosterUser

from modules.app import app
from modules.database import col_banned, col_submitted
from modules.logger import logWrite
from modules.utils import configGet, locale
from classes.enums.submission_types import SubmissionType


@app.on_message(~filters.scheduled & filters.private & filters.photo | filters.video | filters.animation | filters.document)
async def get_submission(app: PosterClient, msg: Message):

    try:

        if col_banned.find_one( {"user": msg.from_user.id} ) is not None:
            return
        
        await app.send_chat_action(msg.chat.id, ChatAction.TYPING)

        user_locale = msg.from_user.language_code
        save_tmp = True
        contents = None

        if PosterUser(msg.from_user.id).is_limited():
            await msg.reply_text(locale("sub_cooldown", "message", locale=user_locale).format(str(configGet("timeout", "submission"))))
            return

        if msg.document is not None:
            if msg.document.mime_type not in configGet("mime_types", "submission"):
                await msg.reply_text(locale("mime_not_allowed", "message", locale=user_locale), quote=True)
                return
            if msg.document.file_size > configGet("file_size", "submission"):
                await msg.reply_text(locale("document_too_large", "message", locale=user_locale).format(str(configGet("file_size", "submission")/1024/1024)), quote=True)
                return
            if msg.document.file_size > configGet("tmp_size", "submission"):
                save_tmp = False
            contents = msg.document.file_id, SubmissionType.DOCUMENT #, msg.document.file_name

        if msg.video is not None:
            if msg.video.file_size > configGet("file_size", "submission"):
                await msg.reply_text(locale("document_too_large", "message", locale=user_locale).format(str(configGet("file_size", "submission")/1024/1024)), quote=True)
                return
            if msg.video.file_size > configGet("tmp_size", "submission"):
                save_tmp = False
            contents = msg.video.file_id, SubmissionType.VIDEO #, msg.video.file_name

        if msg.animation is not None:
            if msg.animation.file_size > configGet("file_size", "submission"):
                await msg.reply_text(locale("document_too_large", "message", locale=user_locale).format(str(configGet("file_size", "submission")/1024/1024)), quote=True)
                return
            if msg.animation.file_size > configGet("tmp_size", "submission"):
                save_tmp = False
            contents = msg.animation.file_id, SubmissionType.ANIMATION #, msg.animation.file_name

        if msg.photo is not None:
            contents = msg.photo.file_id, SubmissionType.PHOTO #, "please_generate"

        if save_tmp is not None:

            if contents is None:
                return

            tmp_id = str(uuid4())
            # filename = tmp_id if contents[1] == "please_generate" else contents[1]
            makedirs(path.join(configGet("data", "locations"), "submissions", tmp_id), exist_ok=True)
            downloaded = await app.download_media(msg, path.join(configGet("data", "locations"), "submissions", tmp_id)+sep)
            inserted = col_submitted.insert_one(
                {
                    "user": msg.from_user.id,
                    "date": datetime.now(),
                    "done": False,
                    "type": contents[1].value,
                    "temp": {
                        "uuid": tmp_id,
                        "file": path.basename(str(downloaded))
                    },
                    "telegram": {
                        "msg_id": msg.id,
                        "file_id": contents[0]
                    },
                    "caption": str(msg.caption) if msg.caption is not None else None
                }
            )
                
        else:

            if contents is None:
                return

            inserted = col_submitted.insert_one(
                {
                    "user": msg.from_user.id,
                    "date": datetime.now(),
                    "done": False,
                    "type": contents[1].value,
                    "temp": {
                        "uuid": None,
                        "file": None
                    },
                    "telegram": {
                        "msg_id": msg.id,
                        "file_id": contents[0]
                    },
                    "caption": str(msg.caption) if msg.caption is not None else None
                }
            )
                
        buttons = [
            [
                InlineKeyboardButton(text=locale("sub_yes", "button", locale=configGet("locale")), callback_data=f"sub_yes_{str(inserted.inserted_id)}")
            ]
        ]

        if msg.caption is not None:
            caption = str(msg.caption)
            buttons[0].append(
                InlineKeyboardButton(text=locale("sub_yes_caption", "button", locale=configGet("locale")), callback_data=f"sub_yes_{str(inserted.inserted_id)}_caption")
            )
            buttons[0].append(
                InlineKeyboardButton(text=locale("sub_no", "button", locale=configGet("locale")), callback_data=f"sub_no_{str(inserted.inserted_id)}")
            )
        else:
            caption = ""
            buttons[0].append(
                InlineKeyboardButton(text=locale("sub_no", "button", locale=configGet("locale")), callback_data=f"sub_no_{str(inserted.inserted_id)}")
            )

        caption += locale("sub_by", "message", locale=locale(configGet("locale")))

        if msg.from_user.first_name is not None:
            caption += f" {msg.from_user.first_name}"
        if msg.from_user.last_name is not None:
            caption += f" {msg.from_user.last_name}"
        if msg.from_user.username is not None:
            caption += f" (@{msg.from_user.username})"
        if msg.from_user.phone_number is not None:
            caption += f" ({msg.from_user.phone_number})"

        if msg.from_user.id in app.admins and configGet("admins", "submission", "require_confirmation") is False:
            try:
                await app.submit_photo(str(inserted.inserted_id))
                await msg.reply_text(locale("sub_yes_auto", "message", locale=user_locale), quote=True)
                await msg.copy(app.owner, caption=caption, disable_notification=True)
                return
            except SubmissionDuplicatesError as exp:
                await msg.reply_text(locale("sub_media_duplicates_list", "message", locale=user_locale).format("\n • ".join(exp.duplicates)), quote=True)
                return
            except Exception as exp:
                await msg.reply_text(format_exc(), quote=True)
                return
        elif msg.from_user.id not in app.admins and configGet("users", "submission", "require_confirmation") is False:
            try:
                await app.submit_photo(str(inserted.inserted_id))
                await msg.reply_text(locale("sub_yes_auto", "message", locale=user_locale), quote=True)
                await msg.copy(app.owner, caption=caption)
                return
            except SubmissionDuplicatesError as exp:
                await msg.reply_text(locale("sub_dup", "message", locale=user_locale), quote=True)
                return
            except Exception as exp:
                await app.send_message(app.owner, locale("sub_error_admin", "message").format(msg.from_user.id, format_exc()))
                await msg.reply_text("sub_error", quote=True)
                return

        if msg.from_user.id not in app.admins:
            buttons += [
                [
                    InlineKeyboardButton(text=locale("sub_block", "button", locale=configGet("locale")), callback_data=f"sub_block_{msg.from_user.id}")
                ]
            ]

            PosterUser(msg.from_user.id).limit()

        if msg.from_user.id != app.owner:
            await msg.reply_text(locale("sub_sent", "message", locale=user_locale), quote=True)

        await msg.copy(app.owner, caption=caption, reply_markup=InlineKeyboardMarkup(buttons))

    except AttributeError:
        logWrite(f"from_user in function get_submission does not seem to contain id")