TelegramPoster/plugins/handlers/submission.py

from datetime import datetime, timezone
from os import makedirs, path, sep
from uuid import uuid4

from pyrogram import filters
from pyrogram.client import Client
from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup, Message

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


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

    try:

        if col_banned.find_one( {"user": msg.from_user.id} ) is not None:
            return

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

        if subLimited(msg.from_user):
            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(tz=timezone.utc),
                    "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(tz=timezone.utc),
                    "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 != configGet("admin"):
            buttons += [
                [
                    InlineKeyboardButton(text=locale("sub_block", "button", locale=configGet("locale")), callback_data=f"sub_block_{msg.from_user.id}")
                ]
                # [
                #     InlineKeyboardButton(text=locale("sub_unblock", "button", locale=configGet("locale")), callback_data=f"sub_unblock_{msg.from_user.id}")
                # ]
            ]

            await msg.reply_text(locale("sub_sent", "message", locale=user_locale), quote=True)
            subLimit(msg.from_user)

        await msg.copy(configGet("admin"), caption=caption, reply_markup=InlineKeyboardMarkup(buttons))

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