telegram_bots/aggregator/client.py
2021-03-11 23:43:54 +03:00

94 lines
3.3 KiB
Python

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()