107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
from datetime import UTC, datetime
|
|
|
|
from pyrogram import Client
|
|
from pyrogram.errors import BadRequest, Forbidden
|
|
from pyrogram.types import User
|
|
|
|
from userbot.modules.avatars import note_avatar
|
|
from userbot.modules.capture.context import CaptureContext
|
|
from userbot.modules.groups.repository import insert_chat_history
|
|
from userbot.modules.jobs.context import JobContext
|
|
from userbot.modules.jobs.registry import register
|
|
from userbot.modules.profiles.parse import snapshot_from_high_level
|
|
from userbot.modules.profiles.repository import write_profile
|
|
from userbot.modules.raw_json import raw_json
|
|
|
|
MEMBER_CAP = 200
|
|
|
|
_MISSING_SENDERS = """
|
|
SELECT DISTINCT sender_id FROM messages
|
|
WHERE account_id = $1 AND chat_id = $2 AND sender_id > 0
|
|
AND NOT EXISTS (
|
|
SELECT 1 FROM peers p WHERE p.account_id = $1 AND p.peer_id = messages.sender_id
|
|
)
|
|
LIMIT 100
|
|
"""
|
|
|
|
|
|
async def _save_user(ctx: CaptureContext, user: User) -> None:
|
|
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
|
|
await write_profile(ctx.pool, ctx.account_id, user.id, fields, raw_json(user))
|
|
if photo_file_id and photo_unique_id:
|
|
await note_avatar(
|
|
ctx.pool, ctx.account_id, user.id, "peer", photo_unique_id, photo_file_id
|
|
)
|
|
|
|
|
|
async def _enrich_chat_meta(client: Client, ctx: CaptureContext, chat_id: int) -> None:
|
|
chat = await client.get_chat(chat_id)
|
|
if chat is None:
|
|
return
|
|
photo = chat.photo
|
|
photo_unique_id = photo.big_photo_unique_id if photo else None
|
|
photo_file_id = photo.big_file_id if photo else None
|
|
await insert_chat_history(
|
|
ctx.pool,
|
|
ctx.account_id,
|
|
chat_id,
|
|
0,
|
|
"meta",
|
|
chat.title,
|
|
photo_unique_id,
|
|
None,
|
|
datetime.now(UTC),
|
|
raw_json(chat),
|
|
)
|
|
if photo_file_id and photo_unique_id:
|
|
await note_avatar(
|
|
ctx.pool, ctx.account_id, chat_id, "chat", photo_unique_id, photo_file_id
|
|
)
|
|
|
|
|
|
async def _enrich_members(client: Client, ctx: CaptureContext, chat_id: int) -> None:
|
|
try:
|
|
async for member in client.get_chat_members(chat_id, limit=MEMBER_CAP):
|
|
if isinstance(member.user, User):
|
|
await _save_user(ctx, member.user)
|
|
except (BadRequest, Forbidden):
|
|
return
|
|
|
|
|
|
async def _fetch_users(client: Client, ids: list[int]) -> list[User]:
|
|
try:
|
|
users = await client.get_users(ids)
|
|
except BadRequest:
|
|
users = []
|
|
for sender_id in ids:
|
|
try:
|
|
users.append(await client.get_users(sender_id))
|
|
except BadRequest:
|
|
continue
|
|
if isinstance(users, User):
|
|
return [users]
|
|
return [user for user in users if isinstance(user, User)]
|
|
|
|
|
|
async def _enrich_senders(client: Client, ctx: CaptureContext, chat_id: int) -> None:
|
|
rows = await ctx.pool.fetch(_MISSING_SENDERS, ctx.account_id, chat_id)
|
|
ids = [row["sender_id"] for row in rows]
|
|
if not ids:
|
|
return
|
|
for user in await _fetch_users(client, ids):
|
|
await _save_user(ctx, user)
|
|
|
|
|
|
@register("enrich_chat")
|
|
async def enrich_chat(ctx: JobContext) -> None:
|
|
client = ctx.client
|
|
if client is None:
|
|
return
|
|
capture = getattr(client, "capture", None)
|
|
if capture is None:
|
|
return
|
|
chat_id = ctx.job.params["chat_id"]
|
|
await _enrich_chat_meta(client, capture, chat_id)
|
|
await _enrich_members(client, capture, chat_id)
|
|
await _enrich_senders(client, capture, chat_id)
|