fix(userbot): presence, story and profile history toggles of the capture policy take effect for each user peer

This commit is contained in:
hh committed 2026-10-19 02:29:19 +00:00
1 parent f32b772ddc
commit 92e5e737d2
10 files changed
+89 -14

No files matched your search

+2 -1
View File
@@ -13,6 +13,7 @@ async def on_user_status(client: PyroClient, user: User) -> None:
ctx = client.capture
if ctx is None or user.status is None:
return
if ctx.resolve_peer(user.id, is_bot=bool(user.is_bot)).presence:
await repository.insert_status(
ctx.pool,
ctx.account_id,
@@ -23,8 +24,8 @@ async def on_user_status(client: PyroClient, user: User) -> None:
user.next_offline_date,
raw_json(user.raw),
)
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)
await ctx.watches.on_status(user.id, is_online=user.status.name.lower() == "online")
handlers = on_user_status.handlers # ty: ignore[unresolved-attribute]
+15 -2
View File
@@ -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(
+8 -3
View File
@@ -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]
+11 -1
View File
@@ -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,
@@ -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
@@ -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
@@ -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,
@@ -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(
@@ -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)
if history:
await conn.execute(_INSERT_HISTORY, account_id, peer_id, *args)
@@ -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