import base64 import json import os import tempfile from io import BytesIO from uuid import uuid4 import imagehash import requests from PIL import Image from django.db import models from telegram import Update, Bot, InputMediaPhoto from telegram.ext import Dispatcher, CallbackContext, MessageHandler, Filters, CommandHandler from jsonrpc import Dispatcher as RPCDispatcher from djconfig import config from bots.models import TelegramBotModuleConfig, BotUser 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 def rpc_post_photo(self, photo, is_base64=False): 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)) im = im.convert('RGB') with tempfile.TemporaryDirectory() as d: fpath = os.path.join(d, '{}.jpg'.format(uuid4())) im.save(fpath) if self.queued: m = bot.send_photo(config.tmp_uploads_chat_id, open(fpath, 'rb')) QueuedItem.objects.create(config=self, type='photo', args=json.dumps([m.photo[-1].file_id]), message_id=m.message_id, image_hash=image_hash) 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 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 self.queued_items.filter(image_hash=i.image_hash, type='photo').count() > 0: update.message.reply_text('Duplicate', 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(message_id=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_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(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) image_hash = models.CharField(max_length=64, null=True, blank=True) processed = models.BooleanField(default=False) 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()