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, Updater class TelegramBot(models.Model): owner = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE) active = models.BooleanField(default=True) title = models.CharField(max_length=32) bot_token = models.CharField(max_length=256) rpc_name = models.CharField(max_length=32, blank=True, null=True) periodic_interval = models.DurationField(blank=True, null=True) periodic_last_run = models.DateTimeField(blank=True, null=True) config_type = models.ForeignKey(ContentType, on_delete=models.CASCADE) config_id = models.PositiveIntegerField() config = GenericForeignKey('config_type', 'config_id') def get_bot(self): return Bot(self.bot_token) def build_updater(self, error_handler): updater = Updater(self.bot_token) updater.dispatcher.add_error_handler(error_handler) updater.bot.get_me() self.config.build_dispatcher(updater.dispatcher) return updater def build_dispatcher(self, error_handler): bot = self.get_bot() bot.get_me() dispatcher = Dispatcher(bot, None, workers=0, use_context=True) dispatcher.add_error_handler(error_handler) 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() self.save() def __str__(self): return f'#{self.pk} {self.title}' class Meta: unique_together = ('config_type', 'config_id') class TelegramBotModuleConfig(models.Model): _bot = GenericRelation(TelegramBot, content_type_field='config_type', object_id_field='config_id') MODULE_NAME = '' @property def bot(self): return self._bot.get() @property def content_type(self): return ContentType.objects.get_for_model(self.__class__) def build_dispatcher(self, dispatcher: Dispatcher): raise NotImplementedError() class Meta: abstract = True class BotUser(models.Model): name = models.CharField(max_length=32) user_id = models.BigIntegerField(db_index=True) def __str__(self): return self.name