captcha.py

from __future__ import annotations

from asyncio import to_thread
from random import sample
from string import ascii_uppercase, digits
from typing import TYPE_CHECKING, NamedTuple

from aiogram.types import BufferedInputFile
from captcha.image import ImageCaptcha


if TYPE_CHECKING:
    from typing import Final, Literal


DEFAULT_FILENAME: Final[str] = "captcha.png"
CHARACTERS: Final[dict[str, str]] = {
    "digits": digits,
    "letters": ascii_uppercase,
    "symbols": digits + ascii_uppercase,
}


class Captcha(NamedTuple):
    image: BufferedInputFile
    characters: str


async def generate_captcha(
    *,
    width: int = 380,
    height: int = 180,
    length: int = 6,
    filename: str = DEFAULT_FILENAME,
    character: Literal["digits", "letters", "symbols"] = "digits",
) -> Captcha:
    characters = CHARACTERS[character]
    length = min(length, len(characters))
    random_characters = "".join(sample(characters, length))
    captcha_creator = ImageCaptcha(width=width, height=height)
    bytes_image = await to_thread(captcha_creator.generate, chars=random_characters)
    buffered_image = BufferedInputFile(file=bytes_image.read(), filename=filename)

    return Captcha(image=buffered_image, characters=random_characters)