telegram_bots/aggregator/client.py

import os
from typing import List

from django.core.files.storage import default_storage
from djconfig import config
from pyrogram import Client
from pyrogram.handlers import MessageHandler, DeletedMessagesHandler
from pyrogram.types import Message as PyrogramMessage
from pyrogram.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: List[PyrogramMessage]):
    for message in 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()