EmojiCaptchaBot/modules/utils.py

from io import BytesIO
from pathlib import Path
from random import randint, sample
from typing import List, Union

from huepaper import generate
from PIL import Image
from pyrogram.enums.chat_member_status import ChatMemberStatus
from pyrogram.types import CallbackQuery, Message

from classes.captcha import Captcha
from classes.pyroclient import PyroClient
from classes.pyrogroup import PyroGroup


def get_captcha_image(emojis: List[str]) -> Captcha:
    emojis_all = sample(emojis, 12)
    emojis_correct = sample(emojis_all, 6)

    output = BytesIO()

    image_options = [
        (randint(400, 490), randint(90, 280), randint(105, 125)),
        (randint(60, 180), randint(180, 240), randint(75, 95)),
        (randint(320, 440), randint(170, 210), randint(35, 45)),
        (randint(150, 240), randint(240, 320), randint(80, 105)),
        (randint(350, 450), randint(280, 380), randint(40, 60)),
        (randint(180, 350), randint(100, 300), randint(45, 65)),
    ]

    base_img = generate(
        width=500,
        height=500,
        hue_max=1.0,
        lum_min=0.3,
        lum_max=0.6,
        sat_min=0.8,
        sat_max=1.0,
    )

    for options, emoji in zip(image_options, emojis_correct):
        width, height, angle = options
        base_img.paste(
            Image.open(Path(f"assets/emojis/{emoji}.png"))
            .resize((120, 120))
            .rotate(angle),
            ((base_img.width - width), (base_img.height - height)),
            Image.open(Path(f"assets/emojis/{emoji}.png"))
            .resize((120, 120))
            .rotate(angle),
        )

    base_img.save(output, format="jpeg")

    return Captcha(output, emojis_all, emojis_correct)


async def is_permitted(
    app: PyroClient,
    group: PyroGroup,
    message: Union[Message, None] = None,
    callback: Union[CallbackQuery, None] = None,
) -> bool:
    """Check if User is an admin or a creator of a group. Alternatively, if the User is actually a group itself.

    ### Args:
        * app (`PyroClient`): Pyrogram Client
        * group (`PyroGroup`): Group
        * message (`Union[Message, None]`, *optional*): Message if the request originates from a command. Defaults to `None`.
        * callback (`Union[CallbackQuery, None]`, *optional*): CallbackQuery if the request originates from a callback. Defaults to `None`.

    ### Returns:
        * `bool`: `True` if permitted and `False` if not. Also `False` if no message or callback provided.
    """
    if message is not None:
        return (
            message.sender_chat is not None and message.sender_chat.id == group.id
        ) or (
            message.from_user is not None
            and (await app.get_chat_member(group.id, message.from_user.id)).status
        ) in [
            ChatMemberStatus.ADMINISTRATOR,
            ChatMemberStatus.OWNER,
        ]

    if callback is not None:
        return (await app.get_chat_member(group.id, callback.from_user.id)).status in [
            ChatMemberStatus.ADMINISTRATOR,
            ChatMemberStatus.OWNER,
        ]

    return False