import base64 import json import logging import os import tempfile import traceback from io import BytesIO from uuid import uuid4 import ffmpeg import imagehash import requests from PIL import Image from django.db import models from django.db.models import Q from telegram import Update, Bot, InputMediaPhoto from telegram.ext import Dispatcher, CallbackContext, MessageHandler, Filters, CommandHandler, CallbackQueryHandler from jsonrpc import Dispatcher as RPCDispatcher from djconfig import config from bots.models import TelegramBotModuleConfig, BotUser from bots.tasks import upload_image_rpc, upload_animation_rpc class ChannelHelperBotModuleConfig(TelegramBotModuleConfig): chat_id = models.CharField(max_length=32) queued = models.BooleanField(default=False) users = models.ManyToManyField(BotUser) send_photo_groups = models.BooleanField(default=False) send_photo_groups_threshold = models.PositiveIntegerField(default=0, help_text='0 = disabled') photo_group_size = models.PositiveIntegerField(default=10) MODULE_NAME = 'Channel helper' rpc_dispatcher = None def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.rpc_dispatcher = RPCDispatcher() self.rpc_dispatcher['post_photo'] = self.rpc_post_photo self.rpc_dispatcher['post_gif'] = self.rpc_post_gif def rpc_post_photo(self, photo, is_base64=False, note=''): config._reload_maybe() bot = self.bot.get_bot() try: if is_base64: f = BytesIO(base64.b64decode(photo)) else: resp = requests.get(photo) resp.raise_for_status() f = BytesIO(resp.content) except: raise RuntimeError('Could not load image') im = Image.open(f) # type: Image.Image image_hash = imagehash.phash(im) if self.queued_items.filter(image_hash=image_hash, type='photo').count() > 0: return 'duplicate' width, height = im.size if width > 2000 or height > 2000: im.thumbnail((2000, 2000), Image.LANCZOS) im = im.convert('RGB') fpath = os.path.join(tempfile.gettempdir(), '{}.jpg'.format(uuid4())) im.save(fpath) qi = QueuedItem.objects.create(config=self, type='photo', args=json.dumps([]), message_id=None, image_hash=image_hash) upload_image_rpc.delay(qi.pk, fpath, note or None) # if self.queued: # m = bot.send_photo(config.tmp_uploads_chat_id, open(fpath, 'rb'), caption=note or None) # else: # bot.send_photo(self.chat_id, open(fpath, 'rb')) # QueuedItem.objects.create(config=self, type='photo', args=json.dumps([]), image_hash=image_hash, processed=True) return True def rpc_post_gif(self, url, note=''): config._reload_maybe() try: fpath_input = os.path.join(tempfile.gettempdir(), '{}.gif'.format(uuid4())) fpath_output = os.path.join(tempfile.gettempdir(), '{}.mp4'.format(uuid4())) resp = requests.get(url) resp.raise_for_status() with open(fpath_input, 'wb') as f: f.write(resp.content) except: raise RuntimeError('Could not load GIF') video_input = ffmpeg \ .input(fpath_input) cmd = video_input \ .output(fpath_output, vf='pad=width=ceil(iw/2)*2:height=ceil(ih/2)*2:x=0:y=0:color=Black', vcodec='libx264', crf='26') logging.info('ffmpeg ' + ' '.join(cmd.get_args())) cmd.run() qi = QueuedItem.objects.create(config=self, type='animation', args=json.dumps([]), message_id=None) upload_animation_rpc.delay(qi.pk, fpath_output, note or None) return True def periodic_task(self, bot: Bot): i = self.queued_items.filter(processed=False).order_by('?').first() # type: QueuedItem if i: i.send(bot) i.processed = True i.save() def handle_message(self, update: Update, ctx: CallbackContext): if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count(): update.effective_message.reply_text('GTFO') return m = update.effective_message bot = ctx.bot i = QueuedItem(config=self, message_id=m.message_id) if hasattr(m, 'audio') and m.audio: a = m.audio i.type = 'audio' i.args = json.dumps([a.file_id, a.duration, a.performer, a.title]) elif hasattr(m, 'document') and m.document: d = m.document i.type = 'document' i.args = json.dumps([d.file_id, d.file_name]) elif hasattr(m, 'photo') and m.photo: p = m.photo i.type = 'photo' i.args = json.dumps([p[-1].file_id]) io = BytesIO() io.name = 'file.jpg' bot.get_file(p[-1]).download(out=io) im = Image.open(io) i.image_hash = imagehash.phash(im) if m.caption and 'force' in m.caption: duplicate = None else: duplicate = self.queued_items.filter(image_hash=i.image_hash, type='photo').first() if duplicate: try: m = update.message.reply_photo(json.loads(duplicate.args)[0], f'Duplicate from {duplicate.datetime}', quote=True) duplicate.message_ids_extra += f'|{m.message_id}' duplicate.save() except: traceback.print_exc() update.message.reply_text('Could not send duplicate original', quote=True) return elif hasattr(m, 'sticker') and m.sticker: s = m.sticker i.type = 'sticker' i.args = json.dumps([s.file_id]) elif hasattr(m, 'video') and m.video: v = m.video i.type = 'video' i.args = json.dumps([v.file_id, v.duration]) elif hasattr(m, 'voice') and m.voice: v = m.voice i.type = 'voice' i.args = json.dumps([v.file_id, v.duration]) elif hasattr(m, 'video_note') and m.video_note: vn = m.video_note i.type = 'video_note' i.args = json.dumps([vn.file_id, vn.duration, vn.length]) elif hasattr(m, 'contact') and m.contact: c = m.contact i.type = 'contact' i.args = json.dumps([c.phone_number, c.first_name, c.last_name]) elif hasattr(m, 'location') and m.location: l = m.location i.type = 'location' i.args = json.dumps([l.latitude, l.longitude]) elif hasattr(m, 'venue') and m.venue: v = m.venue i.type = 'venue' i.args = json.dumps([v.location.latitude, v.location.longitude, v.title, v.address, v.foursquare_id]) elif hasattr(m, 'text') and m.text: i.type = 'message' i.args = json.dumps([m.text_html, 'html']) if self.queued: i.save() else: i.send(bot) def handle_delete(self, update: Update, ctx: CallbackContext): if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count(): update.effective_message.reply_text('GTFO') return reply_to_id = update.effective_message.reply_to_message.message_id try: msg = QueuedItem.objects.get(Q(message_id=reply_to_id) | Q(message_ids_extra__contains=f'{reply_to_id}'), config=self) msg.delete() update.effective_message.reply_text('Deleted') except QueuedItem.DoesNotExist: update.effective_message.reply_text('Not found') def handle_delete_callback(self, update: Update, ctx: CallbackContext): message_id = update.effective_message.message_id try: msg = QueuedItem.objects.get(Q(message_id=message_id) | Q(message_ids_extra__contains=f'{message_id}'), config=self) msg.delete() update.callback_query.answer('Deleted') except QueuedItem.DoesNotExist: update.callback_query.answer('Not found') update.effective_message.edit_caption('Deleted', reply_markup=None) def handle_count(self, update: Update, ctx: CallbackContext): if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count(): update.effective_message.reply_text('GTFO') return update.effective_message.reply_text(f'{QueuedItem.objects.filter(config=self, processed=False).count()} items queued') def build_dispatcher(self, dispatcher: Dispatcher): dispatcher.add_handler(CommandHandler(['delete', 'del', 'remove', 'rem'], self.handle_delete, Filters.reply)) dispatcher.add_handler(CommandHandler(['count'], self.handle_count)) dispatcher.add_handler(CallbackQueryHandler(self.handle_delete_callback, pattern=r'^del (\d+)$')) dispatcher.add_handler(MessageHandler(Filters.private, self.handle_message)) return dispatcher def should_send_photo_group(self): # noinspection PyChainedComparisons return self.send_photo_groups or ( self.send_photo_groups_threshold > 0 and self.queued_items.filter(type='photo', processed=False).count() > self.send_photo_groups_threshold ) def get_additional_images(self, for_id: int): return self.queued_items\ .filter(type='photo', processed=False)\ .exclude(pk=for_id)\ .order_by('?')[:self.photo_group_size - 1] class QueuedItem(models.Model): config = models.ForeignKey(ChannelHelperBotModuleConfig, on_delete=models.CASCADE, related_name='queued_items') type = models.CharField(max_length=12) args = models.TextField() message_id = models.PositiveBigIntegerField(default=None, db_index=True, null=True, blank=True) message_ids_extra = models.TextField(db_index=True, default='') image_hash = models.CharField(max_length=64, null=True, blank=True) processed = models.BooleanField(default=False) datetime = models.DateTimeField(auto_now_add=True) def send(self, bot: Bot): if self.type == 'photo' and self.config.should_send_photo_group(): items = [self] + list(self.config.get_additional_images(self.pk)) bot.send_media_group(self.config.chat_id, [InputMediaPhoto(json.loads(i.args)[0]) for i in items]) for i in items: i.processed = True i.save() else: getattr(bot, 'send_' + self.type)(self.config.chat_id, *json.loads(self.args)) self.processed = True self.save()