openKriemy/commands.py

from io import BytesIO
from struct import pack
from asyncio import create_subprocess_shell
from asyncio.subprocess import PIPE
from datetime import datetime

from tortoise.exceptions import IntegrityError
from telethon.utils import get_display_name, get_peer_id
from telethon.tl.types import MessageEntityCode, MessageEntityPre
from telethon.errors import MessageEmptyError

from aiofiles.os import remove, path
from aiohttp import ClientSession
from emoji import is_emoji
from cairosvg import svg2png

from PIL import Image

from actions import (
    find_action,
    create_action,
    delete_action,
    add_gif,
    add_sticker,
    add_admin,
    delete_admin,
    add_or_update_birthday,
    get_birthdays,
    add_server,
    delete_server,
    add_allowed,
    delete_allowed,
    list_servers,
    get_server_ip,
    is_markov_enabled,
    enable_markov,
    disable_markov,
    set_markov_options,
    get_markov_option,
    markov_say,
    run,
)
from utils import (
    make_temporary_filename,
    make_cache_filename,
    parse_kind,
    get_user_name,
    calculate_age,
    unparse,
)
from rand import roll_dices
from chess0 import IllegalMove, GameOver


class Handler:
    def __init__(self, handler, is_restricted=False, is_public=False):
        self.handler = handler
        self.is_restricted = is_restricted
        self.is_public = is_public


async def newadmin_handler(bot, event, command):
    if command.argc < 1:
        await event.reply("Пожалуйста, укажите пользователя!")

        return

    try:
        target = await bot.get_entity(command.args[0])
    except ValueError:
        await event.reply("Недопустимое имя пользователя!")

        return

    try:
        await add_admin(target)
    except IntegrityError:
        await event.reply("Данный пользователь уже является администратором!")

        return

    await event.reply("Готово!~~")


async def deladmin_handler(bot, event, command):
    if command.argc < 1:
        await event.reply("Пожалуйста, укажите пользователя!")

        return

    try:
        target = await bot.get_entity(command.args[0])
    except ValueError:
        await event.reply("Недопустимое имя пользователя!")

        return

    try:
        await delete_admin(target)
    except IndexError:
        await event.reply("Данный пользователь не является администратором!")

        return

    await event.reply("Готово!~~")


async def newaction_handler(bot, event, command):
    if command.argc < 3:
        await event.reply("Пожалуйста, укажите тип, имя и шаблон действия!")

        return

    try:
        kind = parse_kind(command.args[0])
    except ValueError:
        await event.reply("Неверный тип действия!!!")

        return

    try:
        await create_action(command.args[1], " ".join(command.args[2:]), kind)
    except SyntaxError:
        await event.reply("Недопустимое имя действия!!!")

        return
    except IntegrityError:
        await event.reply("Действие с таким названием уже существует!")

        return

    await event.reply("Действие создано!")


async def delaction_handler(bot, event, command):
    if command.argc < 1:
        await event.reply("Пожалуйста, укажите имя действия!")

        return

    try:
        await delete_action(command.args[0])
    except SyntaxError:
        await event.reply("Недопустимое имя действия!!!")

        return
    except NameError:
        await event.reply("Действия с таким названием не существует!")

        return

    await event.reply("Действие удалено!")


async def addgif_handler(bot, event, command):
    if command.argc < 1:
        await event.reply("Пожалуйста, укажите имя действия!")

        return

    gif = await event.get_reply_message()
    if not gif or not gif.gif:
        await event.reply("Пожалуйста, добавьте GIF!")

        return

    try:
        action = await find_action(command.args[0])

        await add_gif(action, gif.file.id)
    except SyntaxError:
        await event.reply("Недопустимое имя действия!!!")

        return
    except NameError:
        await event.reply("Нет такого действия!")

        return

    await event.reply("Готово!~~")


async def addserver_handler(bot, event, command):
    if command.argc < 2:
        await event.reply("Пожалуйста, укажите имя и адрес сервера!")

        return

    try:
        await add_server(command.args[0], command.args[1])
    except SyntaxError:
        await event.reply("Недопустимое имя сервера!!")

        return
    except ValueError:
        await event.reply("Пожалуйста, введите корректный IPv4-/IPv6-адрес!")

        return
    except IntegrityError:
        await event.reply("Данный сервер уже был добавлен ранее!")

        return

    await event.reply("Готово!~~")


async def delserver_handler(bot, event, command):
    if command.argc < 1:
        await event.reply("Пожалуйста, укажите имя сервера!")

        return

    try:
        await delete_server(command.args[0])
    except SyntaxError:
        await event.reply("Недопустимое имя сервера!!")

        return
    except IndexError:
        await event.reply("Сервер с таким именем не найден!")

        return

    await event.reply("Готово!~~")


async def allow_handler(bot, event, command):
    try:
        await add_allowed(get_peer_id(event.peer_id))
    except IntegrityError:
        await event.reply("Данный чат уже добавлен в список разрешённых!")

        return

    await event.reply("Готово!~~")


async def disallow_handler(bot, event, command):
    try:
        await delete_allowed(get_peer_id(event.peer_id))
    except IndexError:
        await event.reply("Данный чат не найден в списке разрешённых!!")

        return

    await event.reply("Готово!~~")


async def markov_handler(bot, event, command):
    if command.argc < 1:
        await event.reply("Некорректный синтаксис команды!!")

        return

    peer_id = get_peer_id(event.peer_id)

    if command.args[0] == "enable":
        try:
            await enable_markov(peer_id)
        except Exception:
            await event.reply("Ошибка!!!!!")

            return

        await event.reply("Готово!~~")
    elif command.args[0] == "disable":
        try:
            await disable_markov(peer_id)
        except Exception:
            await event.reply("Ошибка!!!!!")

            return

        await event.reply("Готово!~~")
    elif command.args[0] == "is_enabled":
        await event.reply(str(await is_markov_enabled(peer_id)))
    elif command.args[0] == "set" and command.argc == 3:
        try:
            await set_markov_options(
                peer_id, **{command.args[1]: float(command.args[2])}
            )
        except Exception:
            await event.reply("Ошибка!!!!!")

            return

        await event.reply("Готово!~~")
    elif command.args[0] == "get" and command.argc == 2:
        try:
            await event.reply(str(await get_markov_option(peer_id, command.args[1])))
        except Exception:
            await event.reply("Ошибка!!!!!")
    elif command.args[0] == "say":
        if not bot.markov.is_ready:
            await event.reply("Not ready:(")
        else:
            await markov_say(bot, peer_id)
    elif command.args[0] == "reply":
        if not bot.markov.is_ready:
            await event.reply("Not ready:(")
        else:
            await markov_say(bot, peer_id, reply_to=event)
    elif command.args[0] == "is_ready":
        await event.reply(str(bot.markov.is_ready))
    elif command.args[0] == "corpus_size":
        await event.reply(str(len(bot.markov.corpus)))
    elif command.args[0] == "rebuild":
        bot.markov.rebuild()

        await event.reply("Готово!~~")
    elif command.args[0] == "counter":
        await event.reply(str(bot.markov.counter))
    elif command.args[0] == "counter_reset":
        bot.markov.counter = 0

        await event.reply("Готово!~~")
    elif command.args[0] == "reset_corpus":
        bot.markov.corpus = []

        await event.reply("Готово!~~")
    elif command.args[0] == "unready":
        bot.markov.chain = None

        await event.reply("Готово!~~")
    elif command.args[0] == "save":
        bot.markov.save()

        await event.reply("Готово!~~")
    else:
        await event.reply("Некорректный синтаксис команды!!!")


# Very, very, VERY evil code...
async def make_message_shot(bot, message):
    if message.sender is None:
        sender_id = message.peer_id.channel_id

        full_name = await bot.get_entity(sender_id)
        full_name = full_name.title
    else:
        sender_id = message.sender.id

        full_name = get_display_name(message.sender)

    output_path = make_temporary_filename("png")
    avatar_path = await make_cache_filename(sender_id, "png")

    if not await path.isfile(avatar_path):
        await bot.download_profile_photo(sender_id, file=avatar_path, download_big=True)

        # TO-DO: make it better.
        mproc = await create_subprocess_shell(f"mogrify -format png {avatar_path}")
        await mproc.communicate()

    if not await path.isfile(avatar_path):
        avatar_path = "./resources/placeholder.png"

    data = bytes()

    data += pack("I", len(output_path))
    data += bytes(output_path, encoding="ASCII")
    data += pack("I", len(avatar_path))
    data += bytes(avatar_path, encoding="ASCII")

    username = bytes(full_name, encoding="UTF-8")
    data += pack("I", len(username))
    data += username

    data += pack("I", sender_id % 7)

    text = bytes(
        unparse(
            message.raw_text, [entity for entity, _ in message.get_entities_text()]
        ),
        encoding="UTF-8",
    )
    data += pack("I", len(text))
    data += text

    proc = await create_subprocess_shell("./makeshot/makeshot", stdin=PIPE)
    await proc.communicate(input=data)

    pproc = await create_subprocess_shell(f"pngcrush -reduce -ow {output_path}")
    await pproc.communicate()

    return output_path


async def save_handler(bot, event, command):
    message = await event.get_reply_message()
    if not message:
        await event.reply("Пожалуйста, укажите сообщение для сохранения!")

        return

    emoji = "⚡"
    if command.argc >= 1:
        emoji = command.args[0]

        if len(emoji) not in range(1, 6) or not all(map(is_emoji, emoji)):
            await event.reply("Указан некорректный эмодзи!!!")

            return

    path = await make_message_shot(bot, message)

    try:
        file = await add_sticker(bot, path, emoji)

        await bot.send_file(message.peer_id, file=file, reply_to=message)
    finally:
        await remove(path)


async def bday_handler(bot, event, command):
    if command.argc >= 1:
        try:
            date = datetime.strptime(" ".join(command.args), "%d.%m.%Y")
        except ValueError:
            await event.reply(
                "Дата не может быть распознана. Пожалуйста, введите свой день рождения в следующем формате: 01.01.1970 (день, месяц, год)."
            )

            return

        if date >= datetime.now():
            await event.reply("День рождения не может быть в будущем...")

            return

        if await add_or_update_birthday(get_peer_id(event.peer_id), event.sender, date):
            await event.reply("День рождения успешно добавлен!!!")
        else:
            await event.reply("День рождения успешно обновлён!!!")

        return

    birthdays = await get_birthdays(get_peer_id(event.peer_id))
    if not birthdays:
        await event.reply("Пока пусто...")

        return

    birthdays = map(
        lambda birthday: (birthday.user_id, calculate_age(birthday.date)), birthdays
    )
    birthdays = sorted(birthdays, key=lambda birthday: birthday[1].days_until)

    birthdays_list = ""
    for user_id, age in birthdays:
        try:
            entity = await bot.get_entity(user_id)
        except ValueError:
            await bot.get_participants(await event.get_chat())

            try:
                entity = await bot.get_entity(user_id)
            except ValueError:
                continue

        birthdays_list += f"{get_user_name(entity)} ❯ {age.age_now} ➔ {age.age} ❯ {age.date_string} ❯ {age.days_until}\n"

    await event.reply(f"Дни рождения:\n\n{birthdays_list}")


async def vpn_handler(bot, event, command):
    if command.argc < 1:
        await event.reply(
            f"Пожалуйста, укажите имя сервера! Доступные сервера: {await list_servers()}."
        )

        return

    try:
        ip = await get_server_ip(command.args[0])
    except SyntaxError:
        await event.reply("Недопустимое имя сервера!!")

        return
    except IndexError:
        await event.reply("Сервер с таким именем не найден!")

        return

    if event.sender is None:
        sender_id = event.peer_id.channel_id
    else:
        sender_id = event.sender.id

    async with ClientSession() as session:
        try:
            async with session.post(
                f"http://{ip}:9217/api/obtain", data={"user_id": sender_id}
            ) as resp:
                data = await resp.json()

                profile = data["profile"]
        except Exception:
            await event.reply(
                "Произошла ошибка при попытке обращения к API сервера… :("
            )

            return

    try:
        await bot.send_message(
            await bot.get_entity(sender_id),
            f"Ваш файл конфигурации WireGuard (сервер {command.args[0]}):\n\n```{profile}```",
        )
    except Exception:
        await event.reply(
            "Произошла ошибка при отправке файла конфигурации в Ваши личные сообщения… :с"
        )

        return

    await event.reply(
        "Готово!!~~ Файл конфигурации WireGuard отправлен в Ваши личные сообщения!"
    )


async def run_handler(bot, event, command):
    if command.argc < 1:
        async with ClientSession() as session:
            try:
                async with session.get("https://farlands.txlyre.website/langs") as resp:
                    text = await resp.read()
                    text = text.decode("UTF-8")

                    await event.reply(f"Доступные языки:\n`{text}`")
            except Exception:
                await event.reply("Произошла ошибка при попытке обращения к API… :(")

        return

    try:
        _, _, result = await run(command.args_string)

        await event.reply(result)
    except (ValueError, MessageEmptyError):
        await event.reply("<Telegram не смог декодировать текст сообщения>")


async def sylvy_handler(bot, event, command):
    if command.argc < 1:
        await event.reply("Пожалуйста, не оставляйте ввод пустым!")

        return

    async with ClientSession() as session:
        try:
            async with session.post(
                "https://sylvy-engine.txlyre.website/api/compute",
                json={"program": command.args_string, "stringify": True},
            ) as resp:
                data = await resp.json()

                if data["status"] != "ok":
                    await event.reply(f"Ошибка API Sylvy: {data['data']['message']}")

                    return

                result = data["data"]["result"]
        except Exception:
            await event.reply("Произошла ошибка при попытке обращения к API Sylvy… :(")

            return

    if result["status"] == "TIMEOUT":
        await event.reply(
            "Максимальное время исполнения истекло (более тридцати секунд)!!!"
        )

        return
    elif result["status"] != "SUCCESS":
        await event.reply(f"Ошибка исполнения!!!\n\n```{result['stdout']}```")

        return

    if result["plots"]:
        plots = []

        for plot in result["plots"]:
            buffer = BytesIO()
            buffer.name = "plot.png"

            svg2png(bytestring=plot, write_to=buffer)

            buffer.seek(0)

            plots.append(buffer)

        await bot.send_file(
            event.peer_id, file=plots, reply_to=event, force_document=False
        )

        return

    text = ""

    if result["stdout"]:
        text += result["stdout"]
        text += "\n"

    text += result["output"]
    text = text.rstrip()

    await event.reply(
        text, formatting_entities=[MessageEntityCode(offset=0, length=len(text))]
    )


async def say_handler(bot, event, command):
    if not bot.markov.is_ready:
        await event.reply(
            "Генератор текста ещё не готов к использованию. Пожалуйста, попробуйте чуть позже."
        )
    else:
        init_state = None

        if command.argc > 0:
            init_state = command.args_string

        try:
            try:
                await event.delete()
            except Exception:
                pass

            await markov_say(bot, get_peer_id(event.peer_id), init_state=init_state)
        except Exception:
            try:
                await markov_say(bot, get_peer_id(event.peer_id))
            except Exception:
                await event.reply("Ошибка :(")


async def roll_handler(bot, event, command):
    if command.argc < 1:
        await event.reply("Пожалуйста, не оставляйте ввод пустым!")

        return

    try:
        text = await roll_dices(command.args_string)
    except Exception:
        await event.reply("Ошибка :(")

        return

    await event.reply(
        text,
        formatting_entities=[MessageEntityPre(language="", offset=0, length=len(text))],
    )


def chess_render(chess, id, as_list=True):
    board = svg2png(chess.svg(id, shallow=isinstance(id, str)))

    if as_list:
        return [board]

    return board


def chess_game_stats(chess, id):
    if not chess.has_moves(id):
        return "Ходов ещё не сделано."

    text = f"Последние два хода: {chess.moves(id, 2)}\nВсе ходы [{chess.moves_count(id)}]: {chess.moves(id)}"

    if chess.is_checkmate(id):
        text = f"Мат!\n{text}"
    elif chess.is_check(id):
        text = f"Шах!\n{text}"

    return text


def chess_game_over(chess, id, e):
    return [
      chess_render(chess, id, False),
      f"Конец игры: {str(e)}",
      chess_game_stats(chess, id)
    ]


async def chess_start_handler(chess, id):
    await chess.begin(id)

    return chess_render(chess, id)


async def chess_from_handler(chess, id, moves):
    try:
        await chess.begin(id, moves, strict=not isinstance(id, str))
    except GameOver as e:
        return chess_game_over(chess, id, e)
    except IllegalMove as e:
        chess.end(id)

        move = str(e)

        if move:
            return [f"Некорректный ход: {move}"]

        return ["Некорректная последовательность ходов."]

    reply = chess_render(chess, id)

    if chess.is_check(id):
        reply.append("Шах!")

    if isinstance(id, str):
        reply.append(f"Ход {'белых' if chess.turn(id) else 'чёрных'}.")

    return reply


async def chess_stop_handler(chess, id):
    try:
        stats = chess_game_stats(chess, id)
    except KeyError:
        return ["Нет активной игры."]

    chess.end(id)

    return ["Игра завершена.", stats]


async def chess_move_handler(chess, id, move):
    is_group = isinstance(id, str)

    try:
        await chess.move(id, move)

        if not is_group:
            await chess.move(id)
    except KeyError:
        return ["Нет активной игры."]
    except GameOver as e:
        return chess_game_over(chess, id, e)
    except IllegalMove:
        return ["Некорректный ход."]

    reply = chess_render(chess, id)

    if chess.is_check(id):
        reply.append("Шах!")

    return reply


async def chess_undo_handler(chess, id):
    is_group = isinstance(id, str)

    try:
        chess.undo(id)

        if not is_group:
            chess.undo(id)
    except KeyError:
        return ["Нет активной игры."]
    except IndexError:
        return ["Нечего отменять."]
    except GameOver as e:
        return chess_game_over(chess, id, e)

    return ["Последний ход отменён.", svg2png(chess.svg(id))]


async def chess_skip_handler(chess, id):
    is_group = isinstance(id, str)

    try:
        await chess.skip(id)

        if not is_group:
            await chess.move(id)
    except KeyError:
        return ["Нет активной игры."]
    except GameOver as e:
        return chess_game_over(chess, id, e)

    reply = chess_render(chess, id)

    if chess.is_check(id):
        reply.append("Шах!")

    return reply


async def chess_pass_handler(chess, id):
    is_group = isinstance(id, str)

    try:
        await chess.move(id)

        if not is_group:
            await chess.move(id)
    except KeyError:
        return ["Нет активной игры."]
    except GameOver as e:
        return chess_game_over(chess, id, e)

    reply = chess_render(chess, id)

    if chess.is_check(id):
        reply.append("Шах!")

    return reply


async def chess_board_handler(chess, id):     
    try:
        e = chess.game_over(id)
        if e is not None:
            return chess_game_over(chess, id, e)

        reply = chess_render(chess, id)
    except KeyError:
        return ["Нет активной игры."]

    if chess.is_check(id):
        reply.append("Шах!")

    if isinstance(id, str):
        reply.append(f"Ход {'белых' if chess.turn(id) else 'чёрных'}.")

    return reply


async def chess_moves_handler(chess, id):
    try:
        e = chess.game_over(id)
        if e is not None:
            return chess_game_over(chess, id, e)

        text = chess_game_stats(chess, id)
    except KeyError:
        return ["Нет активной игры."]

    if isinstance(id, str):
        text += f"\nХод {'белых' if chess.turn(id) else 'чёрных'}."

    return [text]


async def chess_anim_handler(chess, id, arg):
    count = None

    if arg.strip():
        try:
            count = int(arg)

            if count < 2:
                raise ValueError
        except Exception:
            return ["Некорректное значение аргумента."]

    try:
        frames = chess.animate(id, count=count, shallow=isinstance(id, str))
    except KeyError:
        return ["Нет активной игры."]

    if not frames:
        return ["Ходов ещё не сделано."]

    frames = list(map(lambda f: Image.open(BytesIO(svg2png(f))), frames))
    
    buffer = BytesIO()
    buffer.name = "board.gif"
    frames[0].save(fp=buffer, format="GIF", append_images=frames[1:], save_all=True, duration=500, loop=0)
    buffer.seek(0)

    return [buffer]

CHESS_COMMANDS = {
    "start": (chess_start_handler, "Начать новую игру", 0),
    "from": (chess_from_handler, "Начать новую игру с доской в указанном состоянии", 1, True),
    "end": (chess_stop_handler, "Завершить игру", 0),
    "move": (chess_move_handler, "Сделать ход", 1),
    "undo": (chess_undo_handler, "Отменить ход", 0),
    "skip": (chess_skip_handler, "Пропустить ход", 0),
    "pass": (chess_pass_handler, "Сделать ход вместо вас", 0),
    "board": (chess_board_handler, "Показать доску", 0),
    "moves": (chess_moves_handler, "Показать историю ходов и состояние игры", 0),

    "create": (chess_start_handler, "Создать общую доску", 0),
    "createfrom": (chess_from_handler, "Создать общую доску в указанном состоянии", 1, True),

    "anim": (chess_anim_handler, "Создать анимацию последних N ходов", 0, True),
}

CHESS_ALIASES = {
    "s": "start",
    "f": "from",
    "m": "move",
    "b": "board",
    "u": "undo",
    "s": "skip",
    "p": "pass",
    "ms": "moves",
    "i": "moves",
    "c": "create",
    "cf": "createfrom",
    "a": "anim",
}


async def chess_handler(bot, event, command):
    if not bot.chess:
        await event.reply("Извините, данная функция недоступна....")

        return

    if command.argc < 1:
        buffer = "Доступные субкоманды:\n"

        for command in CHESS_COMMANDS:
            buffer += f"- {command} - {CHESS_COMMANDS[command][1]}.\n"

        await event.reply(buffer.strip())

        return

    name = command.args[0].lower()
    if name in CHESS_ALIASES:
        name = CHESS_ALIASES[name]
    elif name not in CHESS_COMMANDS:
        await event.reply(f"Неизвестная субкоманда: '{name}'.")

        return

    subcommand = CHESS_COMMANDS[name]
    if len(subcommand) == 4 and subcommand[3]:
        if command.argc < subcommand[2]:
            await event.reply(
                f"Некорректное количество аргументов для субкоманды '{name}': требуется как минимум {subcommand[2]}, а получено {command.argc - 1}."
            )

            return

        args = [" ".join(command.args[1:])]
    else:
        if command.argc - 1 != subcommand[2]:
            await event.reply(
                f"Некорректное количество аргументов для субкоманды '{name}': требуется {subcommand[2]}, а получено {command.argc - 1}."
            )

            return

        args = command.args[1:]

    peer_id = str(get_peer_id(event.peer_id))
    result = await subcommand[0](bot.chess, peer_id if name in ("create", "createfrom") or (peer_id in bot.chess.sessions and name not in ("start", "from")) else event.sender.id, *args)

    for reply in result:
        if isinstance(reply, BytesIO):
            await bot.send_file(
                event.peer_id, file=reply, reply_to=event, force_document=True
            )
        elif isinstance(reply, bytes):
            buffer = BytesIO()
            buffer.name = "board.png"
            buffer.write(reply)
            buffer.seek(0)

            await bot.send_file(
                event.peer_id, file=buffer, reply_to=event, force_document=False
            )
        else:
            await event.reply(reply)

    if name in ("move", "skip", "undo", "pass") and peer_id in bot.chess.sessions:
        await event.reply(f"Ход {'белых' if bot.chess.turn(peer_id) else 'чёрных'}.")


_chess_handler = Handler(chess_handler, is_public=True)

COMMANDS = {
    "newadmin": Handler(newadmin_handler, is_restricted=True),
    "deladmin": Handler(deladmin_handler, is_restricted=True),
    "newaction": Handler(newaction_handler, is_restricted=True),
    "delaction": Handler(delaction_handler, is_restricted=True),
    "addgif": Handler(addgif_handler, is_restricted=True),
    "addserver": Handler(addserver_handler, is_restricted=True),
    "delserver": Handler(delserver_handler, is_restricted=True),
    "allow": Handler(allow_handler, is_restricted=True),
    "disallow": Handler(disallow_handler, is_restricted=True),
    "markov": Handler(markov_handler, is_restricted=True),
    "save": Handler(save_handler),
    "bday": Handler(bday_handler),
    # "vpn": Handler(vpn_handler),
    "sylvy": Handler(sylvy_handler, is_public=True),
    "run": Handler(run_handler, is_public=True),
    "say": Handler(say_handler, is_public=True),
    "roll": Handler(roll_handler, is_public=True),
    "chess": _chess_handler, "c": _chess_handler,
}