openKriemy/openkriemy.py

from sys import executable
from asyncio import run, create_subprocess_shell
from asyncio.subprocess import PIPE

from ujson import dumps
from telethon import TelegramClient
from telethon.events import NewMessage
from telethon.utils import resolve_bot_file_id, get_display_name
from tortoise.exceptions import IntegrityError
from aiofiles.os import remove
from emoji import is_emoji

from utils import parse_command, get_link_to_user, make_temporary_filename
from config import config
from db import init_db
from actions import create_action, find_action, delete_action, add_gif, get_random_gif, assign_color, add_sticker

bot = TelegramClient(
  'openkriemy', 
  config.API_ID, 
  config.API_HASH
).start(bot_token=config.API_TOKEN)

# Very, very, VERY evil code...
async def make_message_shot(message):
  proc = await create_subprocess_shell(
    f'{executable} makeshot.py',
    stdin=PIPE
  )

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

  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.username),
    'text': message.text
  }).encode('UTF-8')

  await proc.communicate(input=data) 

  await remove(avatar_path)
 
  return output_path
 
async def handle_command(event, command, argc, args, args_string):
  if command == 'newaction':
    if argc < 2:
      await event.reply('Пожалуйста, укажите имя и шаблон действия!')

      return

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

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

      return

    await event.reply('Действие создано!')
  elif command == 'delaction':
    if argc < 1:
      await event.reply('Пожалуйста, укажите имя действия!')

      return

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

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

      return

    await event.reply('Действие удалено!')
  elif command == 'addgif':
    if 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(args[0])

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

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

      return

    await event.reply('Готово!~~')
  elif command == 'save':
    message = await event.get_reply_message()
    if not message:
      await event.reply('Пожалуйста, укажите сообщение для сохранения!')

      return

    emoji = '⚡'
    if argc >= 1:
      emoji = args[0]

      if not is_emoji(emoji):
        await event.reply('Указан некорректный эмодзи!!!')

        return

    path = await make_message_shot(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)    

@bot.on(NewMessage)
async def on_message(event):
  try:
    command, argc, args, args_string = parse_command(event.text)
  except ValueError:
    return

  if command in (
    'newaction',
    'delaction',
    'addgif',
    'save'
  ):
    await handle_command(event, command, argc, args, args_string)

    return
  
  try:
    action = await find_action(command)
  except SyntaxError:
    return

  if not action:
    return

  reply_to = None
  target = await event.get_reply_message()

  if not target:
    try:
      target = await bot.get_entity(args[0])
    except (ValueError, IndexError):
      await event.reply('Это действие нужно применить на кого-то!')

      return 
  else:
    reply_to = target
    target = target.sender

  await event.delete()
  
  text = action.template.format(**{
    'initiator': get_link_to_user(event.sender),
    'target': get_link_to_user(target)
  })
  
  gif = await get_random_gif(action)

  if gif:
    gif = resolve_bot_file_id(gif.file_id)

  await bot.send_message(
    event.peer_id,
    message=text,
    file=gif,
    reply_to=reply_to
  )
  
with bot:
  bot.loop.run_until_complete(init_db())
  bot.start()
  bot.run_until_disconnected()