104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
import asyncpg
|
|
|
|
from utils.policy import repository
|
|
from utils.policy.models import (
|
|
ChatKind,
|
|
ChatMeta,
|
|
FolderSpec,
|
|
PolicyParent,
|
|
PolicyRecord,
|
|
PolicyRuleView,
|
|
PolicyScopeView,
|
|
PolicySet,
|
|
ScopeType,
|
|
)
|
|
from utils.policy.resolver import resolve_parent
|
|
from utils.read import chats
|
|
from utils.read.models import ChatListItem
|
|
|
|
|
|
def _meta(chat_id: int, chat: ChatListItem | None) -> ChatMeta:
|
|
if chat_id > 0:
|
|
kind = ChatKind.DM
|
|
elif chat is not None and chat.is_broadcast:
|
|
kind = ChatKind.CHANNEL
|
|
else:
|
|
kind = ChatKind.GROUP
|
|
return ChatMeta(
|
|
chat_id=chat_id,
|
|
kind=kind,
|
|
is_bot=bool(chat and chat.is_bot),
|
|
is_contact=chat.is_contact if chat else None,
|
|
)
|
|
|
|
|
|
def _without_chats(policies: PolicySet) -> PolicySet:
|
|
return PolicySet(folder=policies.folder, defaults=policies.defaults)
|
|
|
|
|
|
def _chat_parent(
|
|
chat_id: int,
|
|
chat: ChatListItem | None,
|
|
folders: list[FolderSpec],
|
|
policies: PolicySet,
|
|
) -> PolicyParent:
|
|
return resolve_parent(_meta(chat_id, chat), folders, _without_chats(policies))
|
|
|
|
|
|
async def list_rules(pool: asyncpg.Pool, account_id: int) -> list[PolicyRuleView]:
|
|
records = await repository.list_policies(pool, account_id)
|
|
folders = await repository.list_folders(pool, account_id)
|
|
policies = await repository.load_policy_set(pool, account_id)
|
|
chat_ids = [
|
|
record.scope_id
|
|
for record in records
|
|
if record.scope_type is ScopeType.CHAT and record.scope_id is not None
|
|
]
|
|
known = {
|
|
chat.chat_id: chat for chat in await chats.get_chats(pool, account_id, chat_ids)
|
|
}
|
|
titles = {folder.folder_id: folder.title for folder in folders}
|
|
rules: list[PolicyRuleView] = []
|
|
for record in records:
|
|
rule = PolicyRuleView(**record.model_dump())
|
|
if record.scope_type is ScopeType.FOLDER and record.scope_id is not None:
|
|
rule.title = titles.get(record.scope_id)
|
|
if record.scope_type is ScopeType.CHAT and record.scope_id is not None:
|
|
chat = known.get(record.scope_id)
|
|
rule.title = chat.title if chat else None
|
|
rule.parent = _chat_parent(record.scope_id, chat, folders, policies)
|
|
rules.append(rule)
|
|
return rules
|
|
|
|
|
|
async def _rule(
|
|
pool: asyncpg.Pool, account_id: int, scope_type: ScopeType, scope_id: int
|
|
) -> PolicyRecord | None:
|
|
return await repository.find_policy(pool, account_id, scope_type, scope_id)
|
|
|
|
|
|
async def chat_scope(
|
|
pool: asyncpg.Pool, account_id: int, chat_id: int
|
|
) -> PolicyScopeView:
|
|
folders = await repository.list_folders(pool, account_id)
|
|
policies = await repository.load_policy_set(pool, account_id)
|
|
chat = await chats.get_chat(pool, account_id, chat_id)
|
|
return PolicyScopeView(
|
|
title=chat.title if chat else None,
|
|
rule=await _rule(pool, account_id, ScopeType.CHAT, chat_id),
|
|
parent=_chat_parent(chat_id, chat, folders, policies),
|
|
)
|
|
|
|
|
|
async def folder_scope(
|
|
pool: asyncpg.Pool, account_id: int, folder_id: int
|
|
) -> PolicyScopeView | None:
|
|
folder = await repository.get_folder(pool, account_id, folder_id)
|
|
if folder is None:
|
|
return None
|
|
return PolicyScopeView(
|
|
title=folder.title,
|
|
rule=await _rule(pool, account_id, ScopeType.FOLDER, folder_id),
|
|
parent=None,
|
|
)
|