ForcePilot/backend/package/yuxi/channel/extensions/telegram/security.py

from __future__ import annotations

import logging

logger = logging.getLogger(__name__)

DM_POLICIES = ("pairing", "allowlist", "open", "disabled")
GROUP_POLICIES = ("open", "allowlist", "disabled")


class TelegramSecurity:

    async def check_allowlist(self, peer_id: str, channel_type: str) -> bool:
        return True

    def resolve_dm_policy(self, account: dict | None = None) -> dict:
        if account:
            return self.resolve_dm_policy_for_account(account)
        return {"mode": "pairing", "allow_from": []}

    def resolve_dm_policy_for_account(self, account: dict) -> dict:
        mode = account.get("dm_policy", "pairing")
        return {"mode": mode, "allow_from": account.get("allow_from", [])}

    def resolve_group_policy_for_account(
        self, account: dict, group_id: str | None = None,
    ) -> dict:
        groups = account.get("groups", {})
        if group_id and group_id in groups:
            group_cfg = groups[group_id]
            return {
            "mode": group_cfg.get("groupPolicy", "open"),
            "require_mention": group_cfg.get("requireMention", True),
        }

        wildcard = groups.get("*", {})
        return {"mode": wildcard.get("groupPolicy", account.get("group_policy", "open")),
                "require_mention": wildcard.get("requireMention", True)}

    def resolve_topic_policy(self, account: dict, group_id: str, thread_id: str | None) -> dict | None:
        if not thread_id:
            return None
        groups = account.get("groups", {})
        group_cfg = groups.get(group_id, {})
        topics = group_cfg.get("topics", {})
        if thread_id in topics:
            topic_cfg = topics[thread_id]
            return {
                "agent_id": topic_cfg.get("agentId"),
                "group_policy": topic_cfg.get("groupPolicy"),
                "require_mention": topic_cfg.get("requireMention", True),
                "allow_from": topic_cfg.get("allowFrom", []),
            }
        return None

    def is_allowed_dm(self, account: dict, peer_id: str) -> tuple[bool, str | None]:
        mode = account.get("dm_policy", "pairing")
        if mode == "disabled":
            return False, "DM disabled"

        dms = account.get("dms", {})
        dm_cfg = dms.get(peer_id, {})
        dm_override = dm_cfg.get("dmPolicy")
        effective_mode = dm_override or mode

        if effective_mode == "open":
            return True, None

        allow_from = account.get("allow_from", [])
        normalized_peer = self._normalize_peer(peer_id)

        if effective_mode == "allowlist":
            if self._check_allowlist(allow_from, normalized_peer):
                return True, None
            return False, "not-in-allowlist"

        if effective_mode == "pairing":
            if self._check_allowlist(allow_from, normalized_peer):
                return True, "paired"
            return True, "pairing-required"

        return False, "unknown-policy"

    def is_allowed_group(
        self, account: dict, peer_id: str,
        group_id: str | None = None, is_mentioned: bool = False,
    ) -> tuple[bool, str | None]:
        groups = account.get("groups", {})
        group_cfg = groups.get(group_id, {}) if group_id else {}
        wildcard = groups.get("*", {})

        effective_policy = (
            group_cfg.get("groupPolicy")
            or wildcard.get("groupPolicy")
            or account.get("group_policy", "open")
        )

        if effective_policy == "disabled":
            return False, "group-disabled"

        require_mention = group_cfg.get(
            "requireMention", wildcard.get("requireMention", True),
        )

        if effective_policy == "open":
            if require_mention and not is_mentioned:
                return False, "mention-required"
            return True, None

        if effective_policy == "allowlist":
            allow_from = group_cfg.get("allowFrom", account.get("group_allow_from", []))
            normalized_peer = self._normalize_peer(peer_id)
            if self._check_allowlist(allow_from, normalized_peer):
                return True, None
            return False, "not-in-group-allowlist"

        return False, "unknown-policy"

    @staticmethod
    def _normalize_peer(peer_id: str) -> str:
        for prefix in ("tg:", "telegram:"):
            if peer_id.startswith(prefix):
                return peer_id[len(prefix):]
        return str(peer_id)

    @staticmethod
    def _check_allowlist(allow_from: list[str], peer_id: str) -> bool:
        if "*" in allow_from:
            return True
        normalized = str(peer_id)
        for entry in allow_from:
            entry_normalized = TelegramSecurity._normalize_peer(str(entry))
            if entry_normalized == normalized:
                return True
        return False

    def collect_warnings(
        self, config: dict, account_id: str | None = None, account: dict | None = None
    ) -> list[str]:
        warnings = []
        if not account:
            return warnings

        dm_policy = account.get("dm_policy", "pairing")
        if dm_policy == "open" and not account.get("allow_from"):
            warnings.append("dmPolicy is 'open' without allowFrom — anyone can DM the bot")

        group_policy = account.get("group_policy", "open")
        if group_policy == "open" and not account.get("groups"):
            warnings.append("groupPolicy is 'open' without groups allowlist — bot can be triggered in any group")

        return warnings

    def apply_config_fixes(self, config: dict) -> dict:
        return config