243 lines
10 KiB
Python
243 lines
10 KiB
Python
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
|
|
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))
|
|
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_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)
|
|
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()
|