ForcePilot/backend/package/yuxi/channels/adapters/yuanbao/security.py

from __future__ import annotations

from yuxi.channels.models import ChannelMessage
from yuxi.utils.logging_config import logger


def _resolve_dm_policy(config: dict) -> str:
    if "dm_policy" in config:
        return config["dm_policy"]
    dm = config.get("dm", {})
    if isinstance(dm, dict) and "policy" in dm:
        return dm["policy"]
    return "pairing"


def _resolve_dm_allow_from(config: dict) -> list[str]:
    if "allow_from" in config:
        return config["allow_from"]
    dm = config.get("dm", {})
    if isinstance(dm, dict) and "allowFrom" in dm:
        return dm["allowFrom"]
    return []


async def check_dm_policy(user_id: str, config: dict) -> bool:
    dm_policy = _resolve_dm_policy(config)

    match dm_policy:
        case "open":
            return True
        case "disabled":
            return False
        case "pairing":
            return _is_user_paired(user_id, config)
        case "allowlist":
            allow_from = _resolve_dm_allow_from(config)
            return "*" in allow_from or user_id in allow_from
        case _:
            logger.warning(f"[Yuanbao] Unknown dm_policy '{dm_policy}', denying access for user={user_id}")
            return False


async def check_group_policy(group_open_id: str, user_id: str, config: dict) -> bool:
    groups_config = config.get("groups", {})
    group_cfg = groups_config.get(group_open_id, {})

    if isinstance(group_cfg, dict) and "enabled" in group_cfg and not group_cfg["enabled"]:
        return False

    group_policy = config.get("group_policy", "allowlist")

    match group_policy:
        case "open":
            return True
        case "disabled":
            return False
        case "allowlist":
            group_allow = config.get("group_allow_from", [])
            if "*" in group_allow or user_id in group_allow:
                return True
            per_group_allow = group_cfg.get("allow_from", [])
            return "*" in per_group_allow or user_id in per_group_allow
        case _:
            return False


async def check_mention_required(
    group_open_id: str,
    msg: ChannelMessage,
    config: dict,
    bot_names: list[str] | None = None,
    bot_message_ids: set[str] | None = None,
) -> bool:
    groups_config = config.get("groups", {})
    group_cfg = groups_config.get(group_open_id, {})

    require_mention = group_cfg.get("require_mention")
    if require_mention is None:
        if "group_require_mention" in config:
            require_mention = config["group_require_mention"]
        elif "requireMention" in config:
            require_mention = config["requireMention"]
        else:
            require_mention = True

    if not require_mention:
        return True

    if msg.mentions and msg.mentions.is_bot_mentioned:
        return True

    content = msg.content or ""
    for name in bot_names or []:
        if f"@{name}" in content:
            return True

    reply_to_msg_id = (msg.metadata or {}).get("reply_to_msg_id")
    if reply_to_msg_id and bot_message_ids and reply_to_msg_id in bot_message_ids:
        return True

    return False


def _is_user_paired(user_id: str, config: dict) -> bool:
    paired_users = config.get("paired_users", [])
    return user_id in paired_users