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(CommandHandler(['dump_db'], self.dump_db))
dispatcher.add_handler(CallbackQueryHandler(self.handle_delete_callback, pattern=r'^del (\d+)$'))
dispatcher.add_handler(MessageHandler(Filters.private, self.handle_message))
return dispatcher
def dump_db(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_document(
BytesIO(json.dumps(list(qi.__json__() for qi in QueuedItem.objects.filter(config=self)), indent=2).encode()),
filename='dump.json',
)
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 __json__(self):
return {
'type': self.type,
'args': self.args,
'message_id': self.message_id,
'message_ids_extra': self.message_ids_extra,
'image_hash': self.image_hash,
'processed': self.processed,
'datetime': self.datetime.isoformat(),
}
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()