feat: create message capture policies

backend/src/userbot/folders.py

@@ -0,0 +1,68 @@
+from collections.abc import Iterable
+
+import asyncpg
+from pyrogram import Client, raw, utils
+
+from utils.logging import logger
+from utils.policy.models import FolderSpec
+from utils.policy.repository import replace_folders
+
+
+def _peer_ids(peers: Iterable[raw.base.InputPeer]) -> set[int]:
+    ids: set[int] = set()
+    for peer in peers:
+        try:
+            ids.add(utils.get_peer_id(peer))
+        except (ValueError, AttributeError):
+            continue
+    return ids
+
+
+def _title(raw_title: object) -> str:
+    return getattr(raw_title, "text", None) or str(raw_title)
+
+
+def _parse(raw_filter: raw.base.DialogFilter, order_index: int) -> FolderSpec | None:
+    if isinstance(raw_filter, raw.types.DialogFilterDefault):
+        return None
+    if isinstance(raw_filter, raw.types.DialogFilterChatlist):
+        return FolderSpec(
+            folder_id=raw_filter.id,
+            order_index=order_index,
+            title=_title(raw_filter.title),
+            include_ids=_peer_ids(raw_filter.include_peers),
+            pinned_ids=_peer_ids(raw_filter.pinned_peers),
+            is_chatlist=True,
+        )
+    return FolderSpec(
+        folder_id=raw_filter.id,
+        order_index=order_index,
+        title=_title(raw_filter.title),
+        include_ids=_peer_ids(raw_filter.include_peers),
+        exclude_ids=_peer_ids(raw_filter.exclude_peers),
+        pinned_ids=_peer_ids(raw_filter.pinned_peers),
+        contacts=bool(raw_filter.contacts),
+        non_contacts=bool(raw_filter.non_contacts),
+        groups=bool(raw_filter.groups),
+        broadcasts=bool(raw_filter.broadcasts),
+        bots=bool(raw_filter.bots),
+    )
+
+
+class FolderCache:
+    def __init__(self, client: Client, pool: asyncpg.Pool, account_id: int) -> None:
+        self._client = client
+        self._pool = pool
+        self._account_id = account_id
+        self.folders: list[FolderSpec] = []
+
+    async def refresh(self) -> None:
+        result = await self._client.invoke(raw.functions.messages.GetDialogFilters())
+        specs = [
+            spec
+            for order_index, raw_filter in enumerate(result.filters)
+            if (spec := _parse(raw_filter, order_index)) is not None
+        ]
+        self.folders = specs
+        await replace_folders(self._pool, self._account_id, specs)
+        logger.info(f"[green]Folders cached:[/] {len(specs)}")

backend/src/userbot/handlers/__init__.py

@@ -0,0 +1,3 @@
+from .dialog_filters import dialog_filter_handler
+
+__all__ = ["dialog_filter_handler"]

backend/src/userbot/handlers/dialog_filters.py

@@ -0,0 +1,20 @@
+from pyrogram import Client, raw
+from pyrogram.handlers import RawUpdateHandler
+
+from userbot.folders import FolderCache
+
+_FILTER_UPDATES = (
+    raw.types.UpdateDialogFilter,
+    raw.types.UpdateDialogFilters,
+    raw.types.UpdateDialogFilterOrder,
+)
+
+
+def dialog_filter_handler(cache: FolderCache) -> RawUpdateHandler:
+    async def on_update(
+        _client: Client, update: raw.base.Update, _users: dict, _chats: dict
+    ) -> None:
+        if isinstance(update, _FILTER_UPDATES):
+            await cache.refresh()
+
+    return RawUpdateHandler(on_update)

backend/src/userbot/runner.py

@@ -7,9 +7,14 @@ import asyncpg
 import uvloop
 
 from dependencies.container import container
+from userbot.folders import FolderCache
+from userbot.handlers import dialog_filter_handler
 from userbot.modules import PyroClient
 from utils.env import env
 from utils.logging import logger, setup_logging
+from utils.policy.models import ChatKind, ChatMeta
+from utils.policy.repository import load_policy_set
+from utils.policy.resolver import resolve
 
 setup_logging()
 
@@ -24,6 +29,7 @@ ON CONFLICT (tg_user_id) DO UPDATE SET
     is_active = TRUE,
     raw = EXCLUDED.raw,
     updated_at = now()
+RETURNING account_id
 """
 
 
@@ -34,10 +40,10 @@ def _discover_sessions(sessions_dir: Path) -> list[Path]:
 
 async def _sync_account(
     pool: asyncpg.Pool, client: PyroClient, session_name: str
-) -> None:
+) -> int | None:
     me = client.me
     if not me:
-        return
+        return None
     raw = json.dumps(
         {
             "id": me.id,
@@ -48,10 +54,25 @@ async def _sync_account(
         }
     )
     label = " ".join(filter(None, [me.first_name, me.last_name])) or me.username
-    await pool.execute(
+    account_id = await pool.fetchval(
         _UPSERT_ACCOUNT, me.id, label, me.phone_number, session_name, raw
     )
     logger.info(f"[green]Account synced:[/] {label} ({me.id})")
+    return account_id
+
+
+async def _setup_policy(
+    pool: asyncpg.Pool, client: PyroClient, account_id: int
+) -> None:
+    cache = FolderCache(client, pool, account_id)
+    await cache.refresh()
+    client.add_handler(dialog_filter_handler(cache))
+    if client.me:
+        policies = await load_policy_set(pool, account_id)
+        sample = resolve(
+            ChatMeta(chat_id=client.me.id, kind=ChatKind.DM), cache.folders, policies
+        )
+        logger.info(f"[green]Sample resolve (self DM):[/] {sample.model_dump()}")
 
 
 async def runner() -> None:
@@ -78,10 +99,12 @@ async def runner() -> None:
                 f"{client.me.full_name if client.me else 'unknown'} "
                 f"{client.me.id if client.me else 'unknown'}"
             )
-            await _sync_account(pool, client, session_name)
+            account_id = await _sync_account(pool, client, session_name)
+            if account_id is not None:
+                await _setup_policy(pool, client, account_id)
 
         if clients:
-            logger.info("[green]Userbot running. Idle (no handlers until phase 3).[/]")
+            logger.info("[green]Userbot running.[/]")
             await asyncio.Event().wait()
     finally:
         for client in clients: