Files
beavergram/backend/src/userbot/modules/capture/identity.py
T

111 lines
3.4 KiB
Python

from __future__ import annotations
from typing import TYPE_CHECKING
from userbot.modules.capture.chat_meta import chat_kind
from userbot.modules.profiles.parse import ProfileFields, snapshot_from_high_level
from userbot.modules.profiles.repository import get_peer, write_profile
from utils.policy.models import ChatKind
if TYPE_CHECKING:
import asyncpg
from pyrogram.types import Message
from userbot.modules.capture.context import CaptureContext
class PeerIdentityCache:
def __init__(self) -> None:
self._cache: dict[int, ProfileFields | None] = {}
async def changed(
self, pool: asyncpg.Pool, account_id: int, peer_id: int, fields: ProfileFields
) -> bool:
if peer_id not in self._cache:
self._cache[peer_id] = await get_peer(pool, account_id, peer_id)
if self._cache[peer_id] == fields:
return False
self._cache[peer_id] = fields
return True
class ChatMetaCache:
def __init__(self) -> None:
self._cache: dict[int, tuple[str | None, str | None]] = {}
async def changed(
self,
pool: asyncpg.Pool,
account_id: int,
chat_id: int,
meta: tuple[str | None, str | None],
) -> bool:
from userbot.modules.groups.repository import ( # noqa: PLC0415
get_latest_chat_meta,
)
if chat_id not in self._cache:
self._cache[chat_id] = await get_latest_chat_meta(pool, account_id, chat_id)
if self._cache[chat_id] == meta:
return False
self._cache[chat_id] = meta
return True
async def _capture_peer(message: Message, ctx: CaptureContext) -> None:
user = message.from_user
if user is None:
return
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
if not await ctx.peer_identity.changed(ctx.pool, ctx.account_id, user.id, fields):
return
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user))
if photo_file_id and photo_unique_id:
from userbot.modules.avatars import note_avatar # noqa: PLC0415
await note_avatar(
ctx.pool, ctx.account_id, user.id, "peer", photo_unique_id, photo_file_id
)
async def _capture_chat(message: Message, ctx: CaptureContext) -> None:
chat = message.chat
if (
chat is None
or chat.id is None
or message.date is None
or chat_kind(chat.type) is ChatKind.DM
):
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
meta = (chat.title, photo_unique_id)
if not await ctx.chat_meta.changed(ctx.pool, ctx.account_id, chat.id, meta):
return
from userbot.modules.groups.repository import insert_chat_history # noqa: PLC0415
await insert_chat_history(
ctx.pool,
ctx.account_id,
chat.id,
message.id,
"meta",
chat.title,
photo_unique_id,
None,
message.date,
str(message),
)
if photo_file_id and photo_unique_id:
from userbot.modules.avatars import note_avatar # noqa: PLC0415
await note_avatar(
ctx.pool, ctx.account_id, chat.id, "chat", photo_unique_id, photo_file_id
)
async def capture_identity(message: Message, ctx: CaptureContext) -> None:
await _capture_peer(message, ctx)
await _capture_chat(message, ctx)