Refactor; add admin-restrictions.

actions.py

@@ -7,15 +7,37 @@ from telethon.tl.functions.messages import UploadMediaRequest, GetStickerSetRequ
 from telethon.tl.types import InputStickerSetID, InputStickerSetShortName, InputStickerSetItem, InputMediaUploadedDocument, InputPeerSelf
 from tortoise.expressions import F
 
-from models import Action, Gif, UserColor, StickerPack
+from models import Action, Gif, UserColor, StickerPack, Admin
 from utils import is_valid_name
 from config import config
 
-async def create_action(name, template):
+async def is_admin(bot, user):
+  admin = await bot.get_entity(config.ADMIN)
+
+  if user.id == admin.id:
+    return True
+
+  admin = await Admin.filter(user_id=user.id)
+  if admin:
+    return True
+
+  return False
+
+async def add_admin(user):
+  await Admin(user_id=user.id).save()
+
+async def delete_admin(user):
+  admin = await Admin.filter(user_id=user.id).first()
+  if not admin:
+    raise IndexError
+
+  await admin.delete()
+
+async def create_action(name, template, can_apply_to_self):
   if not is_valid_name(name):
     raise SyntaxError
 
-  await Action(name=name, template=template).save()
+  await Action(name=name, template=template, can_apply_to_self=can_apply_to_self).save()
    
 async def find_action(name):
   if not is_valid_name(name):

commands.py

@@ -0,0 +1,204 @@
+from sys import executable
+from asyncio import create_subprocess_shell
+from asyncio.subprocess import PIPE
+
+from ujson import dumps
+from tortoise.exceptions import IntegrityError
+from telethon.utils import get_display_name
+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
+)
+from utils import make_temporary_filename
+
+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
+
+  can_apply_to_self = command.args[0] == '1'
+
+  try:
+    await create_action(command.args[1], ' '.join(command.args[2:]), can_apply_to_self)
+  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(comman.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')
+
+  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 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 not 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)    
+
+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)
+}

config.yml.sample

@@ -2,6 +2,7 @@ DB_NAME: 'kriemy'
 DB_USER: 'kriemy'
 DB_PASSWORD: '7071'
 
+ADMIN: '@txlyre'
 USER: '@txlyre'
 API_TOKEN: ''
 API_ID: 0

models.py

@@ -1,10 +1,11 @@
 from tortoise.models import Model
-from tortoise.fields import IntField, BigIntField, CharField, TextField, ForeignKeyField
+from tortoise.fields import IntField, BigIntField, CharField, TextField, BooleanField, ForeignKeyField
 
 class Action(Model):
   id = IntField(pk=True)
   name = CharField(max_length=64, unique=True)
   template = TextField()
+  can_apply_to_self = BooleanField(default=True)
 
 class Gif(Model):
   id = IntField(pk=True)
@@ -23,3 +24,7 @@ class StickerPack(Model):
   hash = BigIntField()
 
   stickers_count = IntField(default=0)
+
+class Admin(Model):
+  id = IntField(pk=True)
+  user_id = IntField(unique=True)

openkriemy.py

@@ -1,165 +1,45 @@
-from sys import executable
-from asyncio import run, create_subprocess_shell
-from asyncio.subprocess import PIPE
+from asyncio import run
 
-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 telethon.utils import resolve_bot_file_id
 
-from utils import parse_command, get_link_to_user, make_temporary_filename
+from utils import parse_command, get_link_to_user
 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
+from actions import (
+  find_action, 
+  get_random_gif,
+  is_admin
+)
+from commands import COMMANDS
 
 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)
+    command = parse_command(event.text)
   except ValueError:
     return
 
-  if command in (
-    'newaction',
-    'delaction',
-    'addgif',
-    'save'
-  ):
-    await handle_command(event, command, argc, args, args_string)
+  handler = COMMANDS.get(command.name, None)
+
+  if handler:
+    if handler.is_restricted\
+   and not await is_admin(bot, event.sender):
+      await event.reply('К сожалению, данная команда Вам недоступна.')
+    else:
+      await handler.handler(bot, event, command)
 
     return
-  
+     
   try:
-    action = await find_action(command)
+    action = await find_action(command.name)
   except SyntaxError:
     return
 
@@ -171,7 +51,7 @@ async def on_message(event):
 
   if not target:
     try:
-      target = await bot.get_entity(args[0])
+      target = await bot.get_entity(command.args[0])
     except (ValueError, IndexError):
       await event.reply('Это действие нужно применить на кого-то!')
 
@@ -180,6 +60,12 @@ async def on_message(event):
     reply_to = target
     target = target.sender
 
+  if target.id == event.sender.id\
+ and not action.can_apply_to_self:
+    await event.reply('Данное действие нельзя применять к самому себе...')
+
+    return 
+    
   await event.delete()
   
   text = action.template.format(**{

utils.py

@@ -1,6 +1,12 @@
 from uuid import uuid4
+from collections import namedtuple
 from telethon.utils import get_display_name
 
+Command = namedtuple(
+  'Command', 
+  'name argc args args_string'
+)
+
 def parse_command(text):
   text = text.strip()
 
@@ -22,7 +28,12 @@ def parse_command(text):
   argc = len(args)
   args_string = ' '.join(args)
 
-  return (command, argc, args, args_string)
+  return Command(
+    name=command, 
+    argc=argc, 
+    args=args, 
+    args_string=args_string
+  )
 
 def get_link_to_user(user):
   full_name = get_display_name(user)