telegram_bots/bots/modules/channel_helper.py

205 lines
8.4 KiB
Python
Raw Normal View History

2019-11-12 17:13:56 +00:00
import base64
2019-11-24 16:22:30 +00:00
import json
2019-11-12 17:13:56 +00:00
import os
import tempfile
2023-07-23 00:05:49 +00:00
import traceback
2019-11-12 17:13:56 +00:00
from io import BytesIO
from uuid import uuid4
2022-08-08 01:35:45 +00:00
import imagehash
2019-11-12 17:13:56 +00:00
import requests
from PIL import Image
from django.db import models
2022-08-25 06:44:20 +00:00
from telegram import Update, Bot, InputMediaPhoto
2022-02-22 22:58:23 +00:00
from telegram.ext import Dispatcher, CallbackContext, MessageHandler, Filters, CommandHandler
2019-11-12 17:13:56 +00:00
from jsonrpc import Dispatcher as RPCDispatcher
2019-11-24 16:22:30 +00:00
from djconfig import config
2019-11-12 17:13:56 +00:00
2021-03-15 18:59:34 +00:00
from bots.models import TelegramBotModuleConfig, BotUser
2020-05-05 08:47:40 +00:00
2019-11-12 17:13:56 +00:00
class ChannelHelperBotModuleConfig(TelegramBotModuleConfig):
chat_id = models.CharField(max_length=32)
2019-11-24 16:22:30 +00:00
queued = models.BooleanField(default=False)
2021-03-15 18:59:34 +00:00
users = models.ManyToManyField(BotUser)
2022-08-25 06:44:20 +00:00
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)
2019-11-12 17:13:56 +00:00
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
2023-04-18 21:50:27 +00:00
def rpc_post_photo(self, photo, is_base64=False, note=''):
2019-11-24 16:22:30 +00:00
config._reload_maybe()
bot = self.bot.get_bot()
2019-11-12 17:13:56 +00:00
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
2022-08-08 01:35:45 +00:00
2022-08-08 01:42:13 +00:00
image_hash = imagehash.phash(im)
if self.queued_items.filter(image_hash=image_hash, type='photo').count() > 0:
2022-12-25 14:24:26 +00:00
return 'duplicate'
2022-08-08 01:35:45 +00:00
2019-11-12 17:13:56 +00:00
width, height = im.size
if width > 2000 or height > 2000:
im.thumbnail((2000, 2000))
2019-11-12 17:52:08 +00:00
im = im.convert('RGB')
2019-11-12 17:13:56 +00:00
with tempfile.TemporaryDirectory() as d:
fpath = os.path.join(d, '{}.jpg'.format(uuid4()))
2019-11-12 17:52:08 +00:00
im.save(fpath)
2019-11-24 16:42:28 +00:00
if self.queued:
2023-04-18 21:50:27 +00:00
m = bot.send_photo(config.tmp_uploads_chat_id, open(fpath, 'rb'), caption=note or None)
2022-08-08 01:42:13 +00:00
QueuedItem.objects.create(config=self, type='photo', args=json.dumps([m.photo[-1].file_id]), message_id=m.message_id, image_hash=image_hash)
2019-11-24 16:42:28 +00:00
else:
bot.send_photo(self.chat_id, open(fpath, 'rb'))
2022-08-08 01:42:13 +00:00
QueuedItem.objects.create(config=self, type='photo', args=json.dumps([]), image_hash=image_hash, processed=True)
2019-11-12 17:13:56 +00:00
return True
2019-11-24 16:22:30 +00:00
def periodic_task(self, bot: Bot):
2022-08-08 01:35:45 +00:00
i = self.queued_items.filter(processed=False).order_by('?').first() # type: QueuedItem
2019-11-24 16:22:30 +00:00
if i:
i.send(bot)
2022-08-08 01:35:45 +00:00
i.processed = True
i.save()
2019-11-24 16:22:30 +00:00
2019-11-12 17:13:56 +00:00
def handle_message(self, update: Update, ctx: CallbackContext):
2020-05-05 08:47:40 +00:00
if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count():
update.effective_message.reply_text('GTFO')
return
2019-11-12 17:13:56 +00:00
m = update.effective_message
bot = ctx.bot
2022-02-22 22:58:23 +00:00
i = QueuedItem(config=self, message_id=m.message_id)
2019-11-12 17:13:56 +00:00
if hasattr(m, 'audio') and m.audio:
a = m.audio
2019-11-24 16:22:30 +00:00
i.type = 'audio'
i.args = json.dumps([a.file_id, a.duration, a.performer, a.title])
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'document') and m.document:
d = m.document
2019-11-24 16:22:30 +00:00
i.type = 'document'
i.args = json.dumps([d.file_id, d.file_name])
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'photo') and m.photo:
p = m.photo
2019-11-24 16:22:30 +00:00
i.type = 'photo'
i.args = json.dumps([p[-1].file_id])
2022-08-08 01:35:45 +00:00
io = BytesIO()
io.name = 'file.jpg'
bot.get_file(p[-1]).download(out=io)
im = Image.open(io)
i.image_hash = imagehash.phash(im)
2023-07-23 00:05:49 +00:00
duplicate = self.queued_items.filter(image_hash=i.image_hash, type='photo').first()
if duplicate:
try:
update.message.reply_photo(json.loads(duplicate.args)[0], 'Duplicate', quote=True)
except:
traceback.print_exc()
update.message.reply_text('Could not send duplicate original', quote=True)
2022-08-08 01:35:45 +00:00
return
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'sticker') and m.sticker:
s = m.sticker
2019-11-24 16:22:30 +00:00
i.type = 'sticker'
i.args = json.dumps([s.file_id])
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'video') and m.video:
v = m.video
2019-11-24 16:22:30 +00:00
i.type = 'video'
i.args = json.dumps([v.file_id, v.duration])
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'voice') and m.voice:
v = m.voice
2019-11-24 16:22:30 +00:00
i.type = 'voice'
i.args = json.dumps([v.file_id, v.duration])
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'video_note') and m.video_note:
vn = m.video_note
2019-11-24 16:22:30 +00:00
i.type = 'video_note'
i.args = json.dumps([vn.file_id, vn.duration, vn.length])
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'contact') and m.contact:
c = m.contact
2019-11-24 16:22:30 +00:00
i.type = 'contact'
i.args = json.dumps([c.phone_number, c.first_name, c.last_name])
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'location') and m.location:
l = m.location
2019-11-24 16:22:30 +00:00
i.type = 'location'
i.args = json.dumps([l.latitude, l.longitude])
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'venue') and m.venue:
v = m.venue
2019-11-24 16:22:30 +00:00
i.type = 'venue'
i.args = json.dumps([v.location.latitude, v.location.longitude, v.title, v.address, v.foursquare_id])
2019-11-12 17:13:56 +00:00
elif hasattr(m, 'text') and m.text:
2019-12-04 15:38:14 +00:00
i.type = 'message'
2019-11-24 16:22:30 +00:00
i.args = json.dumps([m.text_html, 'html'])
if self.queued:
i.save()
else:
i.send(bot)
2019-11-12 17:13:56 +00:00
2022-02-22 22:58:23 +00:00
def handle_delete(self, update: Update, ctx: CallbackContext):
2022-04-02 09:42:36 +00:00
if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count():
update.effective_message.reply_text('GTFO')
2022-02-22 22:58:23 +00:00
return
reply_to_id = update.effective_message.reply_to_message.message_id
try:
msg = QueuedItem.objects.get(message_id=reply_to_id, config=self)
2022-02-22 22:58:23 +00:00
msg.delete()
update.effective_message.reply_text('Deleted')
2022-08-04 13:57:01 +00:00
except QueuedItem.DoesNotExist:
update.effective_message.reply_text('Not found')
2022-02-22 22:58:23 +00:00
def handle_count(self, update: Update, ctx: CallbackContext):
2022-04-02 09:42:36 +00:00
if self.users.count() and not self.users.filter(user_id=update.effective_user.id).count():
update.effective_message.reply_text('GTFO')
return
2022-08-26 08:52:05 +00:00
update.effective_message.reply_text(f'{QueuedItem.objects.filter(config=self, processed=False).count()} items queued')
2019-11-12 17:13:56 +00:00
def build_dispatcher(self, dispatcher: Dispatcher):
2022-02-22 22:58:23 +00:00
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))
2019-11-12 17:13:56 +00:00
return dispatcher
2019-11-24 16:22:30 +00:00
2022-08-25 06:44:20 +00:00
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]
2019-11-24 16:22:30 +00:00
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()
2022-02-23 10:13:42 +00:00
message_id = models.PositiveBigIntegerField(default=None, db_index=True, null=True, blank=True)
2022-08-08 01:10:41 +00:00
image_hash = models.CharField(max_length=64, null=True, blank=True)
2022-08-08 01:35:45 +00:00
processed = models.BooleanField(default=False)
2019-11-24 16:22:30 +00:00
def send(self, bot: Bot):
2022-08-25 06:44:20 +00:00
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()