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 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
)
from utils import (
  WORDS_TABLE,
  make_temporary_filename,
  make_cache_filename,
  parse_kind,
  get_user_name,
  get_word_for,
  calculate_age,
  unparse
)

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

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('Готово!~~')

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

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

  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} {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())

      entity = await bot.get_entity(user_id)

    birthdays_list += get_user_name(entity)
    birthdays_list += ' — '
    birthdays_list += age.date_string

    if age.days_until < 1:
      birthdays_list += f' (сегодня исполняется {get_word_for("год", age.age)})\n'
    else:
      birthdays_list += f' (через {get_word_for("день", age.days_until)} исполнится {get_word_for("год", age.age)}; сейчас {get_word_for("год", age.age_now)})\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 sylvy_handler(bot, event, command):
  if command.argc < 1:
    await event.reply('Пожалуйста, не оставляйте ввод пустым!')

    return

  async with ClientSession() as session:
    try:
      async with session.post(
        f'https://sylvy-engine.txlyre.website/api/computeSingle',
        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'] != 'SUCCESS':
     await event.reply(f'Ошибка исполнения!!!\n\n```{result["stdout"]}```')

     return

  if result['plots']:
    plots = []

    for plot in result['plots']:
      buffer = BytesIO()

      svg2png(
        bytestring=plot,
        write_to=buffer
      )

      plots.append(buffer.getvalue())

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

    return

  text = ''

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

  text += result['output']

  await event.reply(f'```{text}```')

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

  'save':      Handler(save_handler),

  'bday':      Handler(bday_handler),

  'vpn':       Handler(vpn_handler),

  'sylvy':     Handler(sylvy_handler),
}