telegram_bots/aggregator/client.py

import os

from django.core.files.storage import default_storage
from djconfig import config
from pyrogram import Client, MessageHandler, DeletedMessagesHandler, \
                     Message as PyrogramMessage, Messages as PyrogramMessages
from pyrogram.api.errors import ChannelPrivate
from pyrogram.session import Session

from aggregator.models import AggregationSource, Message, Chat


Session.notice_displayed = True


def get_client(takeout=False):
    config._reload_maybe()
    if not config.pyrogram_session or not config.pyrogram_app_id or not config.pyrogram_app_hash:
        raise RuntimeError('Pyrogram is not configured')

    session_path = os.path.relpath(default_storage.path(config.pyrogram_session.replace('.session', '')))

    return Client(session_path, config.pyrogram_app_id, config.pyrogram_app_hash, config.pyrogram_app_version or None,
                  config.pyrogram_device_model or None, config.pyrogram_system_version or None, takeout=takeout)


def save_message(client, message: PyrogramMessage):
    if not AggregationSource.objects.filter(chat_id=message.chat.id).exists():
        return
    Message.from_obj(message, client)


def delete_messages(client, messages: PyrogramMessages):
    for message in messages.messages:
        Message.objects.filter(chat__chat_id=message.chat.id, message_id=message.message_id).update(deleted=True)


def collect_new_messages(client, chat):
    # Collecting new messages
    last_message = chat.messages.order_by('-message_id').first()
    if last_message:
        itr = client.iter_history(chat.chat_id, reverse=True, offset_id=last_message.message_id + 1)
    else:
        itr = client.iter_history(chat.chat_id, reverse=True, limit=10)
    for message in itr:
        Message.from_obj(message, client)


def startup_collect(client: Client):
    for chat in Chat.objects.all():
        try:
            client.get_chat(chat.chat_id)
        except ChannelPrivate:
            print('I was banned in chat id="{}"'.format(chat.chat_id))
            continue
        # Collecting edited & deleted messages
        offset = 0
        qs = Message.objects.active_messages().filter(chat__chat_id=chat.chat_id)
        while True:
            lst = qs[200*offset:200*(offset+1)]
            if not lst:
                break
            messages = client.get_messages(chat.chat_id, [m.message_id for m in lst])
            for message in messages.messages:
                m_qs = Message.objects.active_messages() \
                    .filter(chat__chat_id=chat.chat_id, message_id=message.message_id)
                if message.empty:
                     m_qs.update(deleted=True)
                elif message.edit_date:
                    m = m_qs.first()  # type: Message
                    if m and (not m.edit_date or m.edit_date.timestamp() != message.edit_date):
                        Message.from_obj(message, client)
            offset += 1

        collect_new_messages(client, chat)


def start_client():
    app = get_client()

    app.add_handler(MessageHandler(save_message))
    app.add_handler(DeletedMessagesHandler(delete_messages))

    print('Starting pyrogram client...')
    app.start()

    print('Loading updates that happened while I was offline...')
    startup_collect(app)

    print('Idling...')
    app.idle()