From 92e5e737d269ddafab2cdd9c72fab6f5c11b3477 Mon Sep 17 00:00:00 2001 From: h Date: Mon, 19 Oct 2026 02:11:43 +0000 Subject: [PATCH] fix(userbot): presence, story and profile history toggles of the capture policy take effect for each user peer --- backend/src/userbot/handlers/presence.py | 23 ++++++++++--------- backend/src/userbot/handlers/raw/profiles.py | 17 ++++++++++++-- backend/src/userbot/handlers/stories.py | 11 ++++++--- .../src/userbot/modules/capture/context.py | 12 +++++++++- .../src/userbot/modules/capture/identity.py | 9 +++++++- .../modules/jobs/handlers/enrich_chat.py | 9 +++++++- .../modules/jobs/handlers/sync_contacts.py | 11 ++++++++- .../modules/jobs/handlers/sync_dialogs.py | 9 +++++++- .../userbot/modules/profiles/repository.py | 13 ++++++++--- .../src/userbot/modules/profiles/snapshots.py | 11 ++++++++- 10 files changed, 100 insertions(+), 25 deletions(-) diff --git a/backend/src/userbot/handlers/presence.py b/backend/src/userbot/handlers/presence.py index 439d9d0..a41be89 100644 --- a/backend/src/userbot/handlers/presence.py +++ b/backend/src/userbot/handlers/presence.py @@ -13,18 +13,19 @@ async def on_user_status(client: PyroClient, user: User) -> None: ctx = client.capture if ctx is None or user.status is None: return - await repository.insert_status( - ctx.pool, - ctx.account_id, - user.id, - datetime.now(UTC), - user.status.name.lower(), - user.last_online_date, - user.next_offline_date, - raw_json(user.raw), - ) + if ctx.resolve_peer(user.id, is_bot=bool(user.is_bot)).presence: + await repository.insert_status( + ctx.pool, + ctx.account_id, + user.id, + datetime.now(UTC), + user.status.name.lower(), + user.last_online_date, + user.next_offline_date, + raw_json(user.raw), + ) + await notify_bg_event(ctx.pool, "presence", ctx.account_id, chat_id=user.id) await ctx.watches.on_status(user.id, is_online=user.status.name.lower() == "online") - await notify_bg_event(ctx.pool, "presence", ctx.account_id, chat_id=user.id) handlers = on_user_status.handlers # ty: ignore[unresolved-attribute] diff --git a/backend/src/userbot/handlers/raw/profiles.py b/backend/src/userbot/handlers/raw/profiles.py index 2d8fd1f..1e125a0 100644 --- a/backend/src/userbot/handlers/raw/profiles.py +++ b/backend/src/userbot/handlers/raw/profiles.py @@ -25,8 +25,14 @@ async def _handle_user( current = await get_peer(ctx.pool, ctx.account_id, update.user_id) if current == fields: return + is_bot = bool(getattr(raw_user, "bot", False)) await write_profile( - ctx.pool, ctx.account_id, update.user_id, fields, raw_json(raw_user) + ctx.pool, + ctx.account_id, + update.user_id, + fields, + raw_json(raw_user), + history=ctx.resolve_peer(update.user_id, is_bot=is_bot).profile_history, ) changed_photo = current is None or current.photo_unique_id != photo_unique_id if photo_file_id and photo_unique_id and changed_photo: @@ -51,7 +57,14 @@ async def _merge_partial( merged = replace(current or ProfileFields(), **overrides) if merged == current: return - await write_profile(ctx.pool, ctx.account_id, peer_id, merged, raw_str) + await write_profile( + ctx.pool, + ctx.account_id, + peer_id, + merged, + raw_str, + history=ctx.resolve_peer(peer_id).profile_history, + ) async def handle( diff --git a/backend/src/userbot/handlers/stories.py b/backend/src/userbot/handlers/stories.py index fa2eda6..15ff4f1 100644 --- a/backend/src/userbot/handlers/stories.py +++ b/backend/src/userbot/handlers/stories.py @@ -1,14 +1,19 @@ from pyrogram.types import Story from userbot.modules.client import PyroClient -from userbot.modules.stories.service import save_story +from userbot.modules.stories.service import save_story, story_peer_id @PyroClient.on_story() async def on_story(client: PyroClient, story: Story) -> None: - if client.capture is None: + ctx = client.capture + if ctx is None: return - await save_story(client, client.capture, story) + peer_id = story_peer_id(story) + is_bot = bool(story.from_user and story.from_user.is_bot) + if peer_id > 0 and not ctx.resolve_peer(peer_id, is_bot=is_bot).stories: + return + await save_story(client, ctx, story) handlers = on_story.handlers # ty: ignore[unresolved-attribute] diff --git a/backend/src/userbot/modules/capture/context.py b/backend/src/userbot/modules/capture/context.py index e518f47..8f19248 100644 --- a/backend/src/userbot/modules/capture/context.py +++ b/backend/src/userbot/modules/capture/context.py @@ -6,7 +6,7 @@ from userbot.modules.capture.pending import PendingDeletes from userbot.modules.contacts import ContactCache from userbot.modules.folders import FolderCache from userbot.modules.watches import WatchCache -from utils.policy.models import CaptureToggles, ChatMeta, PolicySet +from utils.policy.models import CaptureToggles, ChatKind, ChatMeta, PolicySet from utils.policy.repository import load_policy_set from utils.policy.resolver import resolve from utils.storage import ContentAddressedStorage @@ -38,6 +38,16 @@ class CaptureContext: def resolve(self, chat: ChatMeta) -> CaptureToggles: return resolve(chat, self.folders.folders, self.policies) + def resolve_peer(self, peer_id: int, *, is_bot: bool = False) -> CaptureToggles: + return self.resolve( + ChatMeta( + chat_id=peer_id, + kind=ChatKind.DM, + is_bot=is_bot, + is_contact=peer_id in self.contacts.ids, + ) + ) + async def build_capture_context( client: Client, diff --git a/backend/src/userbot/modules/capture/identity.py b/backend/src/userbot/modules/capture/identity.py index fe1e4c1..909bf0b 100644 --- a/backend/src/userbot/modules/capture/identity.py +++ b/backend/src/userbot/modules/capture/identity.py @@ -58,7 +58,14 @@ async def _capture_peer(message: Message, ctx: CaptureContext) -> None: 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, raw_json(user)) + await write_profile( + ctx.pool, + ctx.account_id, + user.id, + fields, + raw_json(user), + history=ctx.resolve_peer(user.id, is_bot=bool(user.is_bot)).profile_history, + ) if photo_file_id and photo_unique_id: from userbot.modules.avatars import note_avatar # noqa: PLC0415 diff --git a/backend/src/userbot/modules/jobs/handlers/enrich_chat.py b/backend/src/userbot/modules/jobs/handlers/enrich_chat.py index 69008a7..affa944 100644 --- a/backend/src/userbot/modules/jobs/handlers/enrich_chat.py +++ b/backend/src/userbot/modules/jobs/handlers/enrich_chat.py @@ -27,7 +27,14 @@ 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)) + await write_profile( + ctx.pool, + ctx.account_id, + user.id, + fields, + raw_json(user), + history=ctx.resolve_peer(user.id, is_bot=bool(user.is_bot)).profile_history, + ) 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 diff --git a/backend/src/userbot/modules/jobs/handlers/sync_contacts.py b/backend/src/userbot/modules/jobs/handlers/sync_contacts.py index 007883d..02a401b 100644 --- a/backend/src/userbot/modules/jobs/handlers/sync_contacts.py +++ b/backend/src/userbot/modules/jobs/handlers/sync_contacts.py @@ -22,7 +22,16 @@ async def sync_contacts(ctx: JobContext) -> None: if not isinstance(user, User): continue 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)) + await write_profile( + ctx.pool, + ctx.account_id, + user.id, + fields, + raw_json(user), + history=capture.resolve_peer( + user.id, is_bot=bool(user.is_bot) + ).profile_history, + ) if photo_file_id and photo_unique_id: await note_avatar( ctx.pool, diff --git a/backend/src/userbot/modules/jobs/handlers/sync_dialogs.py b/backend/src/userbot/modules/jobs/handlers/sync_dialogs.py index c34c6bc..f009271 100644 --- a/backend/src/userbot/modules/jobs/handlers/sync_dialogs.py +++ b/backend/src/userbot/modules/jobs/handlers/sync_dialogs.py @@ -34,7 +34,14 @@ async def _enrich_users(client: Client, ctx: CaptureContext, ids: list[int]) -> continue 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) + ctx.pool, + ctx.account_id, + user.id, + fields, + raw_json(user), + history=ctx.resolve_peer( + user.id, is_bot=bool(user.is_bot) + ).profile_history, ) if photo_file_id and photo_unique_id: await note_avatar( diff --git a/backend/src/userbot/modules/profiles/repository.py b/backend/src/userbot/modules/profiles/repository.py index 17ce50d..12eb073 100644 --- a/backend/src/userbot/modules/profiles/repository.py +++ b/backend/src/userbot/modules/profiles/repository.py @@ -49,8 +49,14 @@ async def get_peer( ) -async def write_profile( - pool: asyncpg.Pool, account_id: int, peer_id: int, fields: ProfileFields, raw: str +async def write_profile( # noqa: PLR0913 + pool: asyncpg.Pool, + account_id: int, + peer_id: int, + fields: ProfileFields, + raw: str, + *, + history: bool, ) -> None: args = ( fields.first_name, @@ -63,4 +69,5 @@ async def write_profile( ) async with pool.acquire() as conn, conn.transaction(): await conn.execute(_UPSERT_PEER, account_id, peer_id, *args) - await conn.execute(_INSERT_HISTORY, account_id, peer_id, *args) + if history: + await conn.execute(_INSERT_HISTORY, account_id, peer_id, *args) diff --git a/backend/src/userbot/modules/profiles/snapshots.py b/backend/src/userbot/modules/profiles/snapshots.py index feba624..7682493 100644 --- a/backend/src/userbot/modules/profiles/snapshots.py +++ b/backend/src/userbot/modules/profiles/snapshots.py @@ -1,5 +1,6 @@ from datetime import UTC, datetime +from pyrogram.enums import ChatType from pyrogram.types import Chat from userbot.modules.avatars import note_avatar @@ -13,7 +14,15 @@ from userbot.modules.raw_json import raw_json async def save_private(ctx: CaptureContext, chat: Chat) -> bool: chat_id = chat.id or 0 fields, photo_file_id, photo_unique_id = snapshot_from_chat(chat) - await write_profile(ctx.pool, ctx.account_id, chat_id, fields, raw_json(chat)) + is_bot = chat.type is ChatType.BOT + await write_profile( + ctx.pool, + ctx.account_id, + chat_id, + fields, + raw_json(chat), + history=ctx.resolve_peer(chat_id, is_bot=is_bot).profile_history, + ) if photo_file_id and photo_unique_id: await note_avatar( ctx.pool, ctx.account_id, chat_id, "peer", photo_unique_id, photo_file_id