telegram_bots/aggregator/client.py

94 lines
3.3 KiB
Python
Raw Permalink Normal View History

2019-01-27 21:20:09 +00:00
import os
2021-03-11 20:43:54 +00:00
from typing import List
2019-01-27 21:20:09 +00:00
from django.core.files.storage import default_storage
from djconfig import config
2021-03-11 20:43:54 +00:00
from pyrogram import Client
from pyrogram.handlers import MessageHandler, DeletedMessagesHandler
from pyrogram.types import Message as PyrogramMessage
from pyrogram.errors import ChannelPrivate
2019-01-27 21:20:09 +00:00
from pyrogram.session import Session
from aggregator.models import AggregationSource, Message, Chat
Session.notice_displayed = True
2019-01-27 22:07:57 +00:00
def get_client(takeout=False):
2019-01-27 21:20:09 +00:00
config._reload_maybe()
2019-01-27 21:49:17 +00:00
if not config.pyrogram_session or not config.pyrogram_app_id or not config.pyrogram_app_hash:
2019-01-27 21:30:40 +00:00
raise RuntimeError('Pyrogram is not configured')
2019-01-27 21:49:17 +00:00
session_path = os.path.relpath(default_storage.path(config.pyrogram_session.replace('.session', '')))
2019-01-27 21:30:40 +00:00
2019-01-30 17:11:31 +00:00
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)
2019-01-27 21:20:09 +00:00
def save_message(client, message: PyrogramMessage):
if not AggregationSource.objects.filter(chat_id=message.chat.id).exists():
return
Message.from_obj(message, client)
2021-03-11 20:43:54 +00:00
def delete_messages(client, messages: List[PyrogramMessage]):
for message in messages:
2019-01-27 21:20:09 +00:00
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()