implement queue for channel helper

This commit is contained in:
bakatrouble 2019-11-24 19:22:30 +03:00
parent 7725691c1e
commit 9d07c774a3
6 changed files with 113 additions and 19 deletions

View File

@ -3,3 +3,12 @@ from django.apps import AppConfig
class BotsConfig(AppConfig):
name = 'bots'
def ready(self):
self.register_config()
def register_config(self):
import djconfig
from .forms import BotsAppConfigForm
djconfig.register(BotsAppConfigForm)

View File

@ -1,4 +1,6 @@
from django import forms
from django.forms import ModelForm
from djconfig.forms import ConfigForm
from bots.models import TelegramBot
@ -17,3 +19,10 @@ class BotForm(ModelForm):
class Meta:
model = TelegramBot
exclude = 'owner', 'config_type', 'config_id',
class BotsAppConfigForm(ConfigForm):
slug = 'bots'
title = 'Bots'
tmp_uploads_chat_id = forms.CharField(required=True)

View File

@ -0,0 +1,32 @@
# Generated by Django 2.1.5 on 2019-11-24 16:22
from django.db import migrations, models
import django.db.models.deletion
class Migration(migrations.Migration):
dependencies = [
('bots', '0007_auto_20191116_0314'),
]
operations = [
migrations.CreateModel(
name='QueuedItem',
fields=[
('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
('type', models.CharField(max_length=12)),
('args', models.TextField()),
],
),
migrations.AddField(
model_name='channelhelperbotmoduleconfig',
name='queued',
field=models.BooleanField(default=False),
),
migrations.AddField(
model_name='queueditem',
name='config',
field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='queued_items', to='bots.ChannelHelperBotModuleConfig'),
),
]

View File

@ -2,6 +2,7 @@ from django.conf import settings
from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
from django.contrib.contenttypes.models import ContentType
from django.db import models
from django.utils import timezone
from telegram import Bot
from telegram.ext import Dispatcher
@ -31,6 +32,13 @@ class TelegramBot(models.Model):
self.config.build_dispatcher(dispatcher)
return dispatcher
def run_periodic_task(self):
if not hasattr(self.config, 'periodic_task') or not self.periodic_interval or \
(self.periodic_last_run and self.periodic_last_run > timezone.now() - self.periodic_interval):
return
self.config.periodic_task(self.get_bot())
self.periodic_last_run = timezone.now()
def __str__(self):
return f'#{self.pk} {self.title}'

View File

@ -1,4 +1,5 @@
import base64
import json
import os
import tempfile
from io import BytesIO
@ -7,15 +8,17 @@ from uuid import uuid4
import requests
from PIL import Image
from django.db import models
from telegram import Update
from telegram import Update, Bot
from telegram.ext import Dispatcher, CallbackContext, MessageHandler, Filters
from jsonrpc import Dispatcher as RPCDispatcher
from djconfig import config
from bots.models import TelegramBotModuleConfig
class ChannelHelperBotModuleConfig(TelegramBotModuleConfig):
chat_id = models.CharField(max_length=32)
queued = models.BooleanField(default=False)
MODULE_NAME = 'Channel helper'
@ -27,6 +30,8 @@ class ChannelHelperBotModuleConfig(TelegramBotModuleConfig):
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))
@ -44,45 +49,82 @@ class ChannelHelperBotModuleConfig(TelegramBotModuleConfig):
with tempfile.TemporaryDirectory() as d:
fpath = os.path.join(d, '{}.jpg'.format(uuid4()))
im.save(fpath)
self.bot.get_bot().send_photo(self.chat_id, open(fpath, 'rb'))
m = bot.send_photo(config.tmp_uploads_chat_id, open(fpath, 'rb'))
i = QueuedItem(config=self, type='photo', args=json.dumps([m.photo[-1].file_id]))
if self.queued:
i.save()
else:
i.send(bot)
return True
def periodic_task(self, bot: Bot):
i = self.queued_items.order_by('?').first() # type: QueuedItem
if i:
i.send(bot)
i.delete()
def handle_message(self, update: Update, ctx: CallbackContext):
m = update.effective_message
bot = ctx.bot
i = QueuedItem(config=self)
if hasattr(m, 'audio') and m.audio:
a = m.audio
r = bot.send_audio(self.chat_id, a.file_id, a.duration, a.performer, a.title)
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
r = bot.send_document(self.chat_id, d.file_id, d.file_name)
i.type = 'document'
i.args = json.dumps([d.file_id, d.file_name])
elif hasattr(m, 'photo') and m.photo:
p = m.photo
r = bot.send_photo(self.chat_id, p[-1].file_id)
i.type = 'photo'
i.args = json.dumps([p[-1].file_id])
elif hasattr(m, 'sticker') and m.sticker:
s = m.sticker
r = bot.send_sticker(self.chat_id, s.file_id)
i.type = 'sticker'
i.args = json.dumps([s.file_id])
elif hasattr(m, 'video') and m.video:
v = m.video
r = bot.send_video(self.chat_id, v.file_id, v.duration)
i.type = 'video'
i.args = json.dumps([v.file_id, v.duration])
elif hasattr(m, 'voice') and m.voice:
v = m.voice
r = bot.send_voice(self.chat_id, v.file_id, v.duration)
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
r = bot.send_video_note(self.chat_id, vn.file_id, vn.duration, vn.length)
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
r = bot.send_contact(self.chat_id, c.phone_number, c.first_name, c.last_name)
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
r = bot.send_location(self.chat_id, l.latitude, l.longitude)
i.type = 'location'
i.args = json.dumps([l.latitude, l.longitude])
elif hasattr(m, 'venue') and m.venue:
v = m.venue
r = bot.send_venue(self.chat_id, v.location.latitude, v.location.longitude, v.title, v.address, v.foursquare_id)
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:
r = bot.send_message(self.chat_id, m.text_html, 'html')
i.type = 'text'
i.args = json.dumps([m.text_html, 'html'])
if self.queued:
i.save()
else:
i.send(bot)
def build_dispatcher(self, dispatcher: Dispatcher):
dispatcher.add_handler(MessageHandler(Filters.private, self.handle_message))
return dispatcher
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()
def send(self, bot: Bot):
getattr(bot, 'send_' + self.type)(self.config.chat_id, *json.loads(self.args))

View File

@ -1,12 +1,6 @@
import os
import tempfile
from io import BytesIO
import requests
import sentry_sdk
from django.db import models
from djconfig import config
from python_anticaptcha import AnticaptchaClient, ImageToTextTask
from telebot import TeleBot
from telebot.types import InputMediaPhoto
from vk_api import VkApi