TelegramPoster/modules/sender.py

import logging
from datetime import datetime
from os import makedirs, path
from random import choice, sample
from shutil import rmtree
from traceback import format_exc, print_exc
from typing import Union
from uuid import uuid4

import aiofiles
from aiohttp import ClientSession
from libbot.pyrogram.classes import PyroClient
from photosapi_client.errors import UnexpectedStatus
from PIL import Image

from modules.api_client import (
    File,
    PhotoSearch,
    VideoSearch,
    authorize,
    client,
    photo_get,
    photo_patch,
    photo_random,
    video_get,
    video_patch,
    video_random,
)
from modules.database import col_sent, col_submitted

logger = logging.getLogger(__name__)


async def send_content(app: PyroClient, http_session: ClientSession) -> None:
    try:
        try:
            token = await authorize(http_session)
        except ValueError:
            await app.send_message(
                app.owner,
                app._("api_creds_invalid", "message"),
            )
            return

        try:
            funcs = []

            if app.config["posting"]["types"]["photo"]:
                funcs.append((photo_random, photo_get, app.send_photo, photo_patch))

            if app.config["posting"]["types"]["video"]:
                funcs.append((video_random, video_get, app.send_video, video_patch))

            if not funcs:
                raise KeyError(
                    "No media source provided: all seem to be disabled in config"
                )

            if len(funcs) > 1:
                found = False
                for func_iter in sample(funcs, len(funcs)):
                    func = func_iter

                    random_results = (
                        await func_iter[0](
                            album=app.config["posting"]["api"]["album"],
                            caption="queue",
                            client=client,
                            limit=1,
                        )
                    ).results

                    if not random_results:
                        continue

                    media: Union[PhotoSearch, VideoSearch] = random_results[0]

                    try:
                        response: File = await func_iter[1](id=media.id, client=client)
                    except Exception as exp:
                        print_exc()
                        logger.error("Media is invalid: %s", exp)
                        if app.config["reports"]["error"]:
                            await app.send_message(
                                app.owner, f"Media is invalid: {exp}"
                            )
                        return

                    found = True
                    break

                if not found:
                    raise KeyError("No media found")
            else:
                func = funcs[0]
                media: Union[PhotoSearch, VideoSearch] = (
                    await func[0](
                        album=app.config["posting"]["api"]["album"],
                        caption="queue",
                        client=client,
                        limit=1,
                    )
                ).results[0]
                try:
                    response: File = await func[1](id=media.id, client=client)
                except Exception as exp:
                    print_exc()
                    logger.error("Media is invalid: %s", exp)
                    if app.config["reports"]["error"]:
                        await app.send_message(app.owner, f"Media is invalid: {exp}")
                    return

        except (KeyError, AttributeError, TypeError, IndexError):
            logger.info(
                "Could not send content due to queue empty or contains only forbidden extensions"
            )
            if app.config["reports"]["error"]:
                await app.send_message(
                    app.owner,
                    app._("api_queue_empty", "message"),
                )
            return

        except (ValueError, UnexpectedStatus):
            if app.config["reports"]["error"]:
                await app.send_message(
                    app.owner,
                    app._("api_queue_error", "message"),
                )
            return

        tmp_dir = str(uuid4())

        makedirs(path.join(app.config["locations"]["tmp"], tmp_dir), exist_ok=True)

        tmp_path = path.join(tmp_dir, media.filename)

        async with aiofiles.open(
            path.join(app.config["locations"]["tmp"], tmp_path), "wb"
        ) as out_file:
            await out_file.write(response.payload.read())

        logger.info(
            "Candidate %s (%s) is %s bytes big",
            media.filename,
            media.id,
            path.getsize(path.join(app.config["locations"]["tmp"], tmp_path)),
        )

        if (
            path.getsize(path.join(app.config["locations"]["tmp"], tmp_path)) > 5242880
        ) and func[0] is photo_random:
            image = Image.open(path.join(app.config["locations"]["tmp"], tmp_path))
            width, height = image.size
            image = image.resize((int(width / 2), int(height / 2)), Image.ANTIALIAS)
            if tmp_path.lower().endswith(".jpeg") or tmp_path.lower().endswith(".jpg"):
                image.save(
                    path.join(app.config["locations"]["tmp"], tmp_path),
                    "JPEG",
                    optimize=True,
                    quality=50,
                )
            elif tmp_path.lower().endswith(".png"):
                image.save(
                    path.join(app.config["locations"]["tmp"], tmp_path),
                    "PNG",
                    optimize=True,
                    compress_level=8,
                )
            image.close()

        if (
            path.getsize(path.join(app.config["locations"]["tmp"], tmp_path)) > 5242880
        ) and func[0] is photo_random:
            rmtree(
                path.join(app.config["locations"]["tmp"], tmp_dir), ignore_errors=True
            )
            raise BytesWarning

        del response

        submitted = await col_submitted.find_one({"temp.file": media.filename})

        if submitted is not None and submitted["caption"] is not None:
            caption = submitted["caption"].strip()
        else:
            caption = ""

        if (
            submitted is not None
            and app.config["posting"]["submitted_caption"]["enabled"]
            and (
                (submitted["user"] not in app.admins)
                or (
                    app.config["posting"]["submitted_caption"]["ignore_admins"] is False
                )
            )
        ):
            caption = (
                f"{caption}\n\n{app.config['posting']['submitted_caption']['text']}\n"
            )
        else:
            caption = f"{caption}\n\n"

        if app.config["caption"]["enabled"]:
            if app.config["caption"]["link"] is not None:
                caption = f"{caption}[{choice(app.config['caption']['text'])}]({app.config['caption']['link']})"
            else:
                caption = f"{caption}{choice(app.config['caption']['text'])}"
        else:
            caption = caption

        try:
            sent = await func[2](
                app.config["posting"]["channel"],
                path.join(app.config["locations"]["tmp"], tmp_path),
                caption=caption,
                disable_notification=app.config["posting"]["silent"],
            )
        except Exception as exp:
            logger.error(
                "Could not send media %s (%s) due to %s", media.filename, media.id, exp
            )
            if app.config["reports"]["error"]:
                await app.send_message(
                    app.owner,
                    app._("post_exception", "message").format(exp, format_exc()),
                )
            # rmtree(path.join(app.config['locations']['tmp'], tmp_dir), ignore_errors=True)
            return

        await col_sent.insert_one(
            {
                "date": datetime.now(),
                "image": media.id,
                "filename": media.filename,
                "channel": app.config["posting"]["channel"],
                "caption": None
                if (submitted is None or submitted["caption"] is None)
                else submitted["caption"].strip(),
            }
        )

        await func[3](id=media.id, client=client, caption="sent")

        rmtree(path.join(app.config["locations"]["tmp"], tmp_dir), ignore_errors=True)

        logger.info(
            "Sent %s to %s with caption %s and silently %s",
            media.id,
            str(app.config["posting"]["channel"]),
            caption.replace("\n", "%n"),
            str(app.config["posting"]["silent"]),
        )

    except Exception as exp:
        logger.error(
            "Could not send content due to %s. Traceback: %s", exp, format_exc()
        )
        if app.config["reports"]["error"]:
            await app.send_message(
                app.owner,
                app._("post_exception", "message").format(exp, format_exc()),
            )
        try:
            rmtree(
                path.join(app.config["locations"]["tmp"], tmp_dir), ignore_errors=True
            )
        except:
            pass