txlyre 3 years ago
commit
48d9f85b46
13 changed files with 501 additions and 0 deletions
  1. 4 0
      .gitignore
  2. 104 0
      actions.py
  3. 10 0
      config.py
  4. 8 0
      config.yml.sample
  5. 12 0
      db.py
  6. 82 0
      makeshot.py
  7. 25 0
      models.py
  8. 205 0
      openkriemy.py
  9. 5 0
      requirements.txt
  10. BIN
      resources/fonts/OpenSans-Bold.ttf
  11. BIN
      resources/fonts/OpenSans-Regular.ttf
  12. BIN
      resources/tail.png
  13. 46 0
      utils.py

+ 4 - 0
.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+config.yml
+*.session
+

+ 104 - 0
actions.py

@@ -0,0 +1,104 @@
+from random import randint
+
+from tortoise.contrib.postgres.functions import Random
+from telethon.utils import get_input_document
+from telethon.tl.functions.stickers import CreateStickerSetRequest, AddStickerToSetRequest
+from telethon.tl.functions.messages import UploadMediaRequest, GetStickerSetRequest
+from telethon.tl.types import InputStickerSetID, InputStickerSetShortName, InputStickerSetItem, InputMediaUploadedDocument, InputPeerSelf
+from tortoise.expressions import F
+
+from models import Action, Gif, UserColor, StickerPack
+from utils import is_valid_name
+from config import config
+
+async def create_action(name, template):
+  if not is_valid_name(name):
+    raise SyntaxError
+
+  await Action(name=name, template=template).save()
+   
+async def find_action(name):
+  if not is_valid_name(name):
+    raise SyntaxError
+
+  return await Action.filter(name=name).first()
+
+async def delete_action(name):
+  action = await find_action(name)
+  if not action:
+    raise NameError
+
+  gifs = await action.gifs.all()
+  for gif in gifs:
+    await gif.delete()
+
+  await action.delete()
+
+async def add_gif(action, file_id):
+  await Gif(action=action, file_id=file_id).save()
+
+async def get_random_gif(action):
+  return await action.gifs.all().annotate(order=Random()).order_by('order').first()
+
+async def assign_color(username):
+  user_color = await UserColor.filter(username=username).first()
+  if not user_color:
+    user_color = UserColor(username=username, color=randint(0, 7))
+    await user_color.save()
+  
+  return user_color.color
+
+async def create_new_pack(bot, sticker):
+  last_pack = await StickerPack.all().order_by('-id').first()
+  set_id = last_pack.id + 1 if last_pack else 1
+ 
+  user = await bot.get_entity(config.USER)
+  
+  pack = await bot(CreateStickerSetRequest(
+    user_id=user.id,
+    title=f'Messages #{set_id}.',
+    short_name=f'messages{set_id}_by_openkriemybot',
+    stickers=[sticker]
+  ))
+
+  sid = pack.set.id
+  hash = pack.set.access_hash
+
+  await StickerPack(
+    short_name=pack.set.short_name,
+    sid=sid,
+    hash=hash
+  ).save()
+
+  return pack
+
+async def get_current_pack(bot):
+  pack = await StickerPack.all().order_by('-id').first()
+  if not pack or pack.stickers_count == 120:
+    return None
+
+  return pack
+
+async def add_sticker(bot, file, emoji):
+  file = await bot.upload_file(file)
+  file = InputMediaUploadedDocument(file, 'image/png', [])
+  file = await bot(UploadMediaRequest(InputPeerSelf(), file))
+  file = get_input_document(file)
+  sticker = InputStickerSetItem(document=file, emoji=emoji)
+  
+  pack = await get_current_pack(bot)
+  if not pack:
+    pack = await create_new_pack(bot, sticker)
+  else:
+    await StickerPack.filter(id=pack.id).update(stickers_count=F('stickers_count') + 1)
+
+    pack = await bot(AddStickerToSetRequest(
+      stickerset=InputStickerSetID(
+        id=pack.sid,
+        access_hash=pack.hash
+      ), 
+      sticker=sticker
+    ))
+    
+  return get_input_document(pack.documents[-1])
+  

+ 10 - 0
config.py

@@ -0,0 +1,10 @@
+from yaml import load
+from yaml import Loader
+
+class Config:
+  def __init__(self, data):
+    self.__dict__.update(**data)
+
+with open('config.yml') as f:
+  config = load(f, Loader=Loader)
+  config = Config(config)

+ 8 - 0
config.yml.sample

@@ -0,0 +1,8 @@
+DB_NAME: 'kriemy'
+DB_USER: 'kriemy'
+DB_PASSWORD: '7071'
+
+USER: '@txlyre'
+API_TOKEN: ''
+API_ID: 0
+API_HASH: ''

+ 12 - 0
db.py

@@ -0,0 +1,12 @@
+from tortoise import Tortoise, run_async
+
+from config import config
+
+import models
+
+async def init_db():
+  await Tortoise.init(
+    db_url=f'postgres://{config.DB_USER}:{config.DB_PASSWORD}@localhost:5432/{config.DB_NAME}',
+    modules={'models': ['models']}
+  )
+  await Tortoise.generate_schemas() 

+ 82 - 0
makeshot.py

@@ -0,0 +1,82 @@
+from sys import stdin
+from textwrap import fill, shorten
+
+from ujson import loads
+from PIL import Image, ImageDraw, ImageFont
+from emoji import demojize
+
+COLORS = [
+  (0xee, 0x49, 0x28),
+  (0x41, 0xa9, 0x03),
+  (0xe0, 0x96, 0x02),
+  (0x0f, 0x94, 0xed),
+  (0x8f, 0x3b, 0xf7),
+  (0xfc, 0x43, 0x80),
+  (0x00, 0xa1, 0xc4),
+  (0xeb, 0x70, 0x02)
+]
+
+data = loads(stdin.read())
+
+output_path = data['output_path']
+avatar_path = data['avatar_path']
+username = data['username']
+username_color = data['username_color']
+text = data['text']
+
+username = shorten(username, width=15)
+text = fill(text, width=30)
+
+username = demojize(username)
+text = demojize(text)
+
+margin = 24
+hpadding = 18
+vpadding = 14
+avatar_padding = 12
+avatar_dim = 64
+
+font = ImageFont.truetype('./resources/fonts/OpenSans-Regular.ttf', 24)
+username_font = ImageFont.truetype('./resources/fonts/OpenSans-Bold.ttf', 22)
+tail = Image.open('./resources/tail.png')
+mask = Image.new('L', (avatar_dim * 3, avatar_dim * 3), 0)
+draw = ImageDraw.Draw(mask) 
+draw.ellipse((0, 0, avatar_dim * 3, avatar_dim * 3), fill=255)
+mask = mask.resize((avatar_dim, avatar_dim), Image.ANTIALIAS)
+tail_width, tail_height = tail.size
+
+avatar_width = avatar_dim + avatar_padding
+avatar_height = avatar_dim
+
+avatar = Image.open(avatar_path)
+avatar = avatar.resize((avatar_dim, avatar_dim))
+avatar.putalpha(mask)
+
+username_width, username_height = username_font.getsize(username)
+text_width, text_height = font.getsize_multiline(text)
+
+image_width = avatar_width + text_width + username_width + margin*2 + hpadding*2
+image_height = text_height + username_height + margin*2 + vpadding*2
+
+message = Image.new('RGBA', (image_width, image_height))
+message.paste(avatar, (margin, image_height - margin - avatar_height + 2), avatar)
+
+x1 = margin + avatar_dim + avatar_padding
+y1 = margin
+x2 = x1 + text_width + username_width + hpadding*2
+y2 = y1 + text_height + username_height + vpadding*2
+
+draw = ImageDraw.Draw(message)
+draw.rounded_rectangle((x1, y1, x2, y2), fill=(43, 43, 43), radius=7)
+draw.text((margin + avatar_width + hpadding, margin + vpadding/2), username, font=username_font, fill=COLORS[username_color])
+draw.multiline_text((margin + avatar_width + hpadding, margin + username_height + vpadding), text, font=font)
+
+message.paste(tail, (margin + avatar_width - tail_width + 8, image_height - margin - tail_height + 1), tail)
+
+ratio = image_width / image_height
+new_width = 512 if image_width >= image_height else min(int(512 / ratio), 512)
+new_height = 512 if image_height >= image_width else min(int(512 / ratio), 512)
+
+message = message.resize((new_width, new_height))
+
+message.save(output_path)

+ 25 - 0
models.py

@@ -0,0 +1,25 @@
+from tortoise.models import Model
+from tortoise.fields import IntField, BigIntField, CharField, TextField, ForeignKeyField
+
+class Action(Model):
+  id = IntField(pk=True)
+  name = CharField(max_length=64, unique=True)
+  template = TextField()
+
+class Gif(Model):
+  id = IntField(pk=True)
+  action = ForeignKeyField('models.Action', related_name='gifs')
+  file_id = CharField(max_length=64, unique=True)
+
+class UserColor(Model):
+  id = IntField(pk=True)
+  username = CharField(max_length=256, unique=True)
+  color = IntField()
+
+class StickerPack(Model):
+  id = IntField(pk=True)
+  short_name = CharField(max_length=256, unique=True)
+  sid = BigIntField()
+  hash = BigIntField()
+
+  stickers_count = IntField(default=0)

+ 205 - 0
openkriemy.py

@@ -0,0 +1,205 @@
+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)    
+
[email protected](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()

+ 5 - 0
requirements.txt

@@ -0,0 +1,5 @@
+tortoise-orm[asyncpg]
+pyyaml
+aiofiles
+telethon
+emoji

BIN
resources/fonts/OpenSans-Bold.ttf


BIN
resources/fonts/OpenSans-Regular.ttf


BIN
resources/tail.png


+ 46 - 0
utils.py

@@ -0,0 +1,46 @@
+from uuid import uuid4
+from telethon.utils import get_display_name
+
+def parse_command(text):
+  text = text.strip()
+
+  if not text.startswith('/'):
+    raise ValueError
+
+  text = text.split(' ')
+
+  if len(text) < 1:
+    raise ValueError
+
+  command = text[0][1:].lower()
+
+  if '@' in command:
+    command = command.split('@')
+    command = command[0]
+
+  args = text[1:]
+  argc = len(args)
+  args_string = ' '.join(args)
+
+  return (command, argc, args, args_string)
+
+def get_link_to_user(user):
+  full_name = get_display_name(user)
+  if not full_name:
+    full_name = user.username
+  
+  if not full_name:
+    full_name = '?'
+
+  if user.username:
+    return f'[{full_name}](@{user.username})'
+
+  return full_name
+
+def is_valid_name(name):
+  return name.isidentifier()
+
+def make_temporary_filename(ext):
+  uid = uuid4().hex
+
+  return f'tmp_{uid}.{ext}'