telegram_bots/aggregator/client.py

92 lines
3.2 KiB
Python

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