66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
import asyncpg
|
|
from pyrogram import Client
|
|
|
|
from userbot.modules.capture.identity import ChatMetaCache, PeerIdentityCache
|
|
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, ChatKind, ChatMeta, PolicySet
|
|
from utils.policy.repository import load_policy_set
|
|
from utils.policy.resolver import resolve
|
|
from utils.storage import ContentAddressedStorage
|
|
|
|
|
|
class CaptureContext:
|
|
def __init__(
|
|
self,
|
|
account_id: int,
|
|
pool: asyncpg.Pool,
|
|
storage: ContentAddressedStorage,
|
|
folders: FolderCache,
|
|
contacts: ContactCache,
|
|
) -> None:
|
|
self.account_id = account_id
|
|
self.pool = pool
|
|
self.storage = storage
|
|
self.folders = folders
|
|
self.contacts = contacts
|
|
self.watches = WatchCache(pool, account_id)
|
|
self.peer_identity = PeerIdentityCache()
|
|
self.chat_meta = ChatMetaCache()
|
|
self.pending_deletes = PendingDeletes()
|
|
self.policies = PolicySet()
|
|
|
|
async def reload_policies(self) -> None:
|
|
self.policies = await load_policy_set(self.pool, self.account_id)
|
|
|
|
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,
|
|
pool: asyncpg.Pool,
|
|
storage: ContentAddressedStorage,
|
|
account_id: int,
|
|
) -> CaptureContext:
|
|
folders = FolderCache(client, pool, account_id)
|
|
await folders.refresh()
|
|
contacts = ContactCache(client)
|
|
await contacts.refresh()
|
|
ctx = CaptureContext(account_id, pool, storage, folders, contacts)
|
|
await ctx.reload_policies()
|
|
await ctx.watches.refresh()
|
|
return ctx
|