openKriemy/commands.py

from sys import executable
from asyncio import create_subprocess_shell
from asyncio.subprocess import PIPE
from datetime import datetime, timedelta

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,
  assign_color,
  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_kind

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(
    f'{executable} makeshot.py',
    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)

  data = dumps({
    'output_path': output_path,
    'avatar_path': avatar_path,
    'username': full_name if full_name else message.sender.username,
    'username_color': await assign_color(message.sender.id if message.sender else message.peer_id.channel_id),
    'text': message.text
  }).encode('UTF-8')

  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 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 = sorted(birthdays, key=lambda birthday: birthday.date)

  birthdays_list = ''
  for birthday in birthdays:   
    now = datetime.now().date()
    birthday_date = birthday.date.replace(year=now.year)

    if now > birthday_date:
      birthday_date = birthday.date.replace(year=now.year + 1)

    delta = birthday_date - now
    age = (birthday_date - birthday.date).days // 365
    age_now = age - 1

    birthdays_list += get_user_name(await bot.get_entity(birthday.user_id))
    birthdays_list += ' — '
    birthdays_list += birthday.date.strftime('%d.%m.%Y')
    birthdays_list += f' (через {delta.days} {WORDS_TABLE["день"][get_word_kind(delta.days)]} исполнится {age} {WORDS_TABLE["год"][get_word_kind(age)]}; сейчас {age_now} {WORDS_TABLE["год"][get_word_kind(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)
}