TelegramPoster/classes/poster_client.py

from os import path, remove, sep
from shutil import rmtree
from typing import Union
from pyrogram.client import Client
from pyrogram.types import Message
from classes.exceptions import SubmissionDuplicatesError, SubmissionUnavailableError
from modules.api_client import upload_pic
from modules.database import col_submitted
from bson import ObjectId
from modules.logger import logWrite

from modules.utils import configGet

class PosterClient(Client):

    def __init__(self, name: str, **kwargs): # type: ignore
        super().__init__(name, **kwargs)
        self.owner = configGet("owner")
        self.admins = configGet("admins")+[configGet("owner")]

    async def submit_photo(self, id: str) -> Union[Message, None]:
        
        db_entry = col_submitted.find_one({"_id": ObjectId(id)})
        submission = None

        if db_entry is None:
            raise SubmissionUnavailableError()
        else:
            if db_entry["temp"]["uuid"] is not None:
                if not path.exists(path.join(configGet("data", "locations"), "submissions", db_entry["temp"]["uuid"], db_entry["temp"]["file"])):
                    raise SubmissionUnavailableError()
                else:
                    filepath = path.join(configGet("data", "locations"), "submissions", db_entry["temp"]["uuid"], db_entry["temp"]["file"])
                    try:
                        submission = await self.get_messages(db_entry["user"], db_entry["telegram"]["msg_id"])
                    except:
                        pass
            else:
                try:
                    submission = await self.get_messages(db_entry["user"], db_entry["telegram"]["msg_id"])
                    filepath = await self.download_media(submission, file_name=configGet("tmp", "locations")+sep)
                except:
                    raise SubmissionUnavailableError()
                
        response = await upload_pic(str(filepath))

        if len(response[1]) > 0:
            raise SubmissionDuplicatesError(str(filepath), response[1])
        
        col_submitted.find_one_and_update({"_id": ObjectId(id)}, {"$set": {"done": True}})

        try:
            if db_entry["temp"]["uuid"] is not None:
                rmtree(path.join(configGet("data", "locations"), "submissions", db_entry["temp"]["uuid"]), ignore_errors=True)
            else:
                remove(str(filepath))
        except (FileNotFoundError, NotADirectoryError):
            logWrite(f"Could not delete '{filepath}' on submission accepted", debug=True)

        return submission

    async def ban_user(self, id: int) -> None:
        pass

    async def unban_user(self, id: int) -> None:
        pass