openKriemy/commands.py

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

from ujson import dumps
from tortoise.exceptions import IntegrityError
from telethon.utils import get_display_name, get_peer_id
from telethon.tl.types import MessageEntityCode

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

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,
)


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:
            await event.reply("Ошибка!!!!!")

            return

        await event.reply("Готово!~~")
    elif command.args[0] == "disable":
        try:
            await disable_markov(peer_id)
        except:
            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:
            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:
            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:
            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:
        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:
                await event.reply("Произошла ошибка при попытке обращения к API… :(")

        return

    await event.reply(await run(command.args_string))


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:
            await event.reply("Произошла ошибка при попытке обращения к API Sylvy… :(")

            return

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

        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:
                pass

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


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),
}