From 2d315593334da84a902a32ac729bbeee06d350e8 Mon Sep 17 00:00:00 2001 From: udf Date: Mon, 22 Apr 2019 22:32:36 +0200 Subject: [PATCH] Initial commit --- .gitignore | 4 +++ bot.py | 84 +++++++++++++++++++++++++++++++++++++++++++++++++++ downloader.py | 63 ++++++++++++++++++++++++++++++++++++++ proxy.py | 11 +++++++ rescaler.py | 45 +++++++++++++++++++++++++++ uploader.py | 63 ++++++++++++++++++++++++++++++++++++++ util.py | 58 +++++++++++++++++++++++++++++++++++ 7 files changed, 328 insertions(+) create mode 100644 .gitignore create mode 100644 bot.py create mode 100644 downloader.py create mode 100644 proxy.py create mode 100644 rescaler.py create mode 100644 uploader.py create mode 100644 util.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ef0c48b --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +.vscode +.idea +__pycache__ +*.session diff --git a/bot.py b/bot.py new file mode 100644 index 0000000..d1a583a --- /dev/null +++ b/bot.py @@ -0,0 +1,84 @@ +import asyncio +import logging +import importlib +import uuid +import traceback + +from telethon import TelegramClient +from telethon import events +from telethon.tl.types import DocumentAttributeSticker +from telethon.tl.functions.messages import GetStickerSetRequest + +from util import StatusMessage, Job, find_instance +import proxy + + +proxy.JOB_MODULES = JOB_MODULES = {} + + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger('main') +proxy.client = client = TelegramClient("bot", 6, "eb06d4abfb49dc3eeb1aeb98ae0f581e").start() +proxy.me = me = asyncio.get_event_loop().run_until_complete(client.get_me()) + + +@client.on(events.NewMessage(pattern=r'/start')) +async def on_start(event): + await event.respond('Send a sticker.') + + +@client.on(events.NewMessage) +async def on_message(event): + if not event.message.sticker: + return + sticker = event.message.sticker + sticker_attrib = find_instance(sticker.attributes, DocumentAttributeSticker) + if not sticker_attrib.stickerset: + await event.reply('That sticker is not part of a pack') + return + + sticker_set = await client(GetStickerSetRequest(sticker_attrib.stickerset)) + + status = StatusMessage(await event.reply('Pending')) + job = Job( + id=uuid.uuid1(), + owner=await event.get_input_sender(), + event=event, + status=status, + sticker_set=sticker_set + ) + logger.info( + '[%s] User %s requested sticker set %s', + job.id, + (await event.get_input_sender()).user_id, + sticker_set.set.short_name + ) + + await status.update('Waiting for download slot.') + await JOB_MODULES['downloader'].queue.put(job) + await status.update('Queued for download.') + + +async def job_runner(mod): + while 1: + try: + await mod.run_job(await mod.queue.get()) + except Exception as e: + logger.log( + 'Exception on job runner for %s\n%s', + mod.__name__, + traceback.format_exc() + ) + + +def load_handler_module(name): + proxy.logger = logging.getLogger(name) + mod = importlib.import_module(name) + asyncio.ensure_future(job_runner(mod)) + return mod + + +for name in ('downloader', 'rescaler', 'uploader'): + JOB_MODULES[name] = load_handler_module(name) + +client.run_until_disconnected() diff --git a/downloader.py b/downloader.py new file mode 100644 index 0000000..87b0feb --- /dev/null +++ b/downloader.py @@ -0,0 +1,63 @@ +""" +Handles downloading sticker packs to memory +Based on my gist here: https://gist.github.com/udf/e4e3dbb2e831c8b580d8fddd312714f7 +""" +import asyncio +from io import BytesIO +from collections import defaultdict + +from util import Job, Sticker +from proxy import client, logger, JOB_MODULES + + +queue = asyncio.Queue(5) + + +async def run_job(job: Job): + logger.info(f'[{job.id}] Running download job') + await job.status.update('Starting download...') + + sticker_set = job.sticker_set + + # Sticker emojis are retrieved as a mapping of + # : + # So we need to build a mapping of : + # Thanks, Durov + emojis = defaultdict(str) + for pack in sticker_set.packs: + for document_id in pack.documents: + emojis[document_id] += pack.emoticon + + pending_tasks = [] + stickers = [] + for i, document in enumerate(sticker_set.documents): + file = BytesIO() + task = asyncio.ensure_future( + client.download_media(document, file=file) + ) + pending_tasks.append(task) + stickers.append(Sticker(file, emojis[document.id])) + + await job.status.update('Downloading...') + await asyncio.wait(pending_tasks) + + errors = '' + for i, task in reversed(tuple(enumerate(pending_tasks))): + exception = task.exception() + if not exception: + continue + errors += f'#{i}: {exception.__class__.__name__}\n' + del stickers[i] + + if errors: + await job.event.reply( + f'Some errors occured when downloading these stickers:\n{errors}' + "\nI'll continue to process the pack as normal, but the stickers " + "which caused these errors won't be present in the new pack" + ) + + job.stickers = stickers + logger.info(f'[{job.id}] Finished running download job') + await job.status.update('Waiting for rescale slot...') + await JOB_MODULES['rescaler'].queue.put(job) + await job.status.update('Queued for rescaling.') diff --git a/proxy.py b/proxy.py new file mode 100644 index 0000000..e72de93 --- /dev/null +++ b/proxy.py @@ -0,0 +1,11 @@ +from logging import Logger +from typing import Dict +from types import ModuleType + +from telethon import TelegramClient +from telethon.tl import types + +client: TelegramClient = None +logger: Logger = None +JOB_MODULES: Dict[str, ModuleType] = {} +me: types.User = None diff --git a/rescaler.py b/rescaler.py new file mode 100644 index 0000000..58509b8 --- /dev/null +++ b/rescaler.py @@ -0,0 +1,45 @@ +""" +Handles spawning an executor to liquid rescale the stickers +""" +import asyncio +import concurrent.futures +from io import BytesIO + +from wand.image import Image + +from util import Job, Sticker +from proxy import logger, JOB_MODULES + +queue = asyncio.Queue(3) +executor = concurrent.futures.ThreadPoolExecutor() + + +def crunch(sticker: Sticker): + sticker.file.seek(0) + with Image(file=sticker.file) as i: + i.resize(width=i.width * 2, height=i.height * 2, filter='sinc') + i.liquid_rescale(i.width // 2, i.height // 2, delta_x=1, rigidity=5) + sticker.file = BytesIO() + i.save(file=sticker.file) + + +def do_crunch(loop, job: Job): + futures = [executor.submit(crunch, sticker) for sticker in job.stickers] + for i, future in enumerate(concurrent.futures.as_completed(futures)): + asyncio.run_coroutine_threadsafe( + job.status.update(f'Rescaled {i+1}/{len(futures)}', important=False), + loop + ).result() + + +async def run_job(job: Job): + logger.info(f'[{job.id}] Running rescale job') + await job.status.update('Crunching...') + + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, lambda: do_crunch(loop, job)) + + logger.info(f'[{job.id}] Finished running rescale job') + await job.status.update('Waiting for upload slot...') + await JOB_MODULES['uploader'].queue.put(job) + await job.status.update('Queued for upload.') diff --git a/uploader.py b/uploader.py new file mode 100644 index 0000000..ed1172c --- /dev/null +++ b/uploader.py @@ -0,0 +1,63 @@ +""" +Handles uploading the rescaled stickers and figuring out an unused url +""" +import asyncio +import re + +from telethon import utils +from telethon.errors import ShortnameOccupyFailedError, PackShortNameOccupiedError +from telethon.tl import types +from telethon.tl.functions.messages import UploadMediaRequest +from telethon.tl.functions.stickers import CreateStickerSetRequest + +from util import Job, get_pack_url, get_rand_letters +from proxy import client, logger, me + + +queue = asyncio.Queue(5) +re_by_us = re.compile(f'(?i)_[a-z]+_by_{me.username}$') +re_by_bot = re.compile(r'(?i)_by_\w+bot$') + + +async def run_job(job: Job): + logger.info(f'[{job.id}] Running upload job') + await job.status.update('Uploading...') + stickers = [] + for i, sticker in enumerate(job.stickers): + sticker.file.seek(0) + file = await client.upload_file(sticker.file, part_size_kb=512) + file = types.InputMediaUploadedDocument(file, 'image/png', []) + media = await client(UploadMediaRequest('me', file)) + stickers.append(types.InputStickerSetItem( + document=utils.get_input_document(media), + emoji=sticker.emoji + )) + await job.status.update( + f'Uploaded {i + 1}/{len(job.stickers)}', + important=False + ) + + # Create + id_len = 2 + title = f'Distorted {job.sticker_set.set.title}'[:64] + original_short_name = re_by_us.sub('', job.sticker_set.set.short_name) + original_short_name = re_by_bot.sub('', original_short_name) + while 1: + try: + suffix = f'_{get_rand_letters(id_len)}_by_{me.username}' + short_name = f'{original_short_name}'[:62 - len(suffix)] + suffix + logger.info('Trying to create pack %s', short_name) + await client(CreateStickerSetRequest( + user_id=await job.event.get_input_sender(), + title=title, + short_name=short_name, + stickers=stickers + )) + break + except (ShortnameOccupyFailedError, PackShortNameOccupiedError): + id_len = min(7, id_len + 1) + continue + + logger.info(f'[{job.id}] Finished running upload job') + await job.event.reply('Done! ' + get_pack_url(short_name)) + await job.status.message.delete() diff --git a/util.py b/util.py new file mode 100644 index 0000000..f20e86c --- /dev/null +++ b/util.py @@ -0,0 +1,58 @@ +import time +import random +import string +from io import BytesIO +from dataclasses import dataclass +from typing import List + +from telethon import events +from telethon.tl import types + + +class StatusMessage: + def __init__(self, message): + self.message = message + self.last_edit = time.time() + + async def update(self, text, important=True): + current_time = time.time() + if not important and current_time - self.last_edit < 5: + return + self.last_edit = current_time + await self.message.edit(text) + + +@dataclass +class Sticker: + file: BytesIO + emoji: str + + +@dataclass +class Job: + # Common + id: str + owner: types.InputUser + event: events.NewMessage + status: StatusMessage + + # Downloader + sticker_set: types.messages.StickerSet + + # Resizer, Uploader + stickers: List[Sticker] = None + + +def find_instance(items, class_or_tuple): + for item in items: + if isinstance(item, class_or_tuple): + return item + return None + + +def get_pack_url(short_name): + return f'https://t.me/addstickers/{short_name}' + + +def get_rand_letters(k): + return ''.join(random.choices(string.ascii_uppercase, k=k))