openKriemy/commands.py

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
from emoji import is_emoji

from actions import (
  find_action,
  create_action,
  delete_action,
  add_gif,
  add_sticker,
  add_admin,
  delete_admin,
  add_or_update_birthday,
  get_birthdays
)
from utils import (
  WORDS_TABLE,
  make_temporary_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('Готово!~~')

# Very, very, VERY evil code...
async def make_message_shot(bot, message):
  proc = await create_subprocess_shell(
    './makeshot/makeshot',
    stdin=PIPE
  )

  output_path = make_temporary_filename('png')
  avatar_path = make_temporary_filename('png')

  if message.sender is None:
    await bot.download_profile_photo(message.peer_id.channel_id, file=avatar_path)

    full_name = await bot.get_entity(message.peer_id.channel_id)
    full_name = full_name.title
  else:
    await bot.download_profile_photo(message.sender, file=avatar_path)

    full_name = get_display_name(message.sender)

  # 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 if full_name else message.sender.username, encoding='UTF-8')
  data += pack('I', len(username))
  data += username

  data += pack('I', (message.sender.id if message.sender else message.peer_id.channel_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

  await proc.communicate(input=data) 

  await remove(avatar_path)
 
  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:   
    birthdays_list += get_user_name(await bot.get_entity(user_id))
    birthdays_list += ' — '
    birthdays_list += age.date_string
    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}')

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

  'save':      Handler(save_handler),

  'bday':      Handler(bday_handler)
}