ForcePilot/backend/package/yuxi/channels/adapters/googlechat/policy.py

from __future__ import annotations

from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any

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


class DmPolicy(StrEnum):
    OPEN = "open"
    PAIRING = "pairing"
    ALLOWLIST = "allowlist"
    DISABLED = "disabled"


class GroupPolicy(StrEnum):
    OPEN = "open"
    ALLOWLIST = "allowlist"
    DISABLED = "disabled"


DM_POLICY_VALUES = frozenset(p.value for p in DmPolicy)
GROUP_POLICY_VALUES = frozenset(p.value for p in GroupPolicy)


@dataclass
class PerGroupConfig:
    require_mention: bool = False
    enabled: bool = True
    users: list[str] = field(default_factory=list)
    system_prompt: str = ""


@dataclass
class GoogleChatPolicy:
    dm_policy: DmPolicy = DmPolicy.OPEN
    group_policy: GroupPolicy = GroupPolicy.OPEN
    allow_from: list[str] = field(default_factory=list)
    group_allow_from: list[str] = field(default_factory=list)
    require_mention: bool = False
    allow_bots: bool = False
    bot_user: str = ""
    groups: dict[str, PerGroupConfig] = field(default_factory=dict)
    dangerously_allow_name_matching: bool = False

    @classmethod
    def from_config(cls, config: dict[str, Any] | None) -> GoogleChatPolicy:
        if not config:
            return cls()

        dm_raw = str(config.get("dm_policy", config.get("dmPolicy", "open"))).strip().lower()
        dm_policy = DmPolicy(dm_raw) if dm_raw in DM_POLICY_VALUES else DmPolicy.OPEN

        group_raw = str(config.get("group_policy", config.get("groupPolicy", "open"))).strip().lower()
        group_policy = GroupPolicy(group_raw) if group_raw in GROUP_POLICY_VALUES else GroupPolicy.OPEN

        allow_from = _normalize_list(config.get("allowFrom", config.get("allow_from", [])))
        group_allow_from = _normalize_list(config.get("groupAllowFrom", config.get("group_allow_from", [])))
        require_mention = bool(config.get("require_mention", config.get("requireMention", False)))
        allow_bots = bool(config.get("allow_bots", config.get("allowBots", False)))
        bot_user = str(config.get("bot_user", config.get("botUser", ""))).strip()
        groups = _parse_per_group_config(config.get("groups", {}))
        dangerously_allow_name_matching = bool(
            config.get("dangerously_allow_name_matching", config.get("dangerouslyAllowNameMatching", False))
        )

        return cls(
            dm_policy=dm_policy,
            group_policy=group_policy,
            allow_from=allow_from,
            group_allow_from=group_allow_from,
            require_mention=require_mention,
            allow_bots=allow_bots,
            bot_user=bot_user,
            groups=groups,
            dangerously_allow_name_matching=dangerously_allow_name_matching,
        )

    def check_dm_access(self, user_id: str) -> bool:
        if self.dm_policy == DmPolicy.DISABLED:
            return False
        if self.dm_policy == DmPolicy.OPEN:
            return True
        if self.dm_policy == DmPolicy.PAIRING:
            return True
        if self.dm_policy == DmPolicy.ALLOWLIST:
            return self._match_allowlist(user_id, self.allow_from)
        return False

    def check_group_access(self, space_id: str) -> bool:
        per_group = self._resolve_per_group(space_id)
        if per_group is not None and not per_group.enabled:
            return False
        if self.group_policy == GroupPolicy.DISABLED:
            return False
        if self.group_policy == GroupPolicy.OPEN:
            return True
        if self.group_policy == GroupPolicy.ALLOWLIST:
            return self._match_allowlist(space_id, self.group_allow_from)
        return False

    def check_inbound(self, channel_msg: ChannelMessage) -> bool:
        chat_type = channel_msg.chat_type

        if chat_type == ChatType.DIRECT:
            user_id = channel_msg.identity.channel_user_id
            return self.check_dm_access(user_id)

        space_id = channel_msg.metadata.get("space_name", "")
        if not self.check_group_access(space_id):
            return False

        per_group = self._resolve_per_group(space_id)

        if per_group is not None and per_group.users:
            user_id = channel_msg.identity.channel_user_id
            if user_id not in per_group.users:
                logger.debug(f"Group message filtered: user {user_id} not in per-group users whitelist")
                return False

        require_mention = self.require_mention
        if per_group is not None:
            require_mention = per_group.require_mention

        if require_mention and chat_type != ChatType.DIRECT:
            mentions = channel_msg.mentions
            if mentions and not mentions.is_bot_mentioned:
                logger.debug("Group message filtered: bot not mentioned")
                return False

        return True

    def get_per_group_system_prompt(self, space_id: str) -> str:
        per_group = self._resolve_per_group(space_id)
        if per_group and per_group.system_prompt:
            return per_group.system_prompt
        return ""

    def _resolve_per_group(self, space_id: str) -> PerGroupConfig | None:
        if not space_id or not self.groups:
            return None
        return self.groups.get(space_id)

    def _match_allowlist(self, target: str, allowlist: list[str]) -> bool:
        if not allowlist:
            return False
        if "*" in allowlist:
            return True
        if target in allowlist:
            return True
        if self.dangerously_allow_name_matching and "@" in target:
            email_local = target.split("@")[0].lower()
            for entry in allowlist:
                if "@" in entry and entry.split("@")[0].lower() == email_local:
                    logger.warning(
                        f"Dangerously matched target '{target}' to allowlist entry '{entry}' via email local-part"
                    )
                    return True
        return False


def _normalize_list(raw: Any) -> list[str]:
    if isinstance(raw, str):
        return [x.strip() for x in raw.split(",") if x.strip()]
    if isinstance(raw, (list, tuple)):
        return [str(x).strip() for x in raw if x]
    return []


def _parse_per_group_config(groups_raw: Any) -> dict[str, PerGroupConfig]:
    if not isinstance(groups_raw, dict):
        return {}
    result: dict[str, PerGroupConfig] = {}
    for space_name, cfg in groups_raw.items():
        if not isinstance(cfg, dict):
            continue
        result[space_name] = PerGroupConfig(
            require_mention=bool(cfg.get("require_mention", cfg.get("requireMention", False))),
            enabled=bool(cfg.get("enabled", True)),
            users=_normalize_list(cfg.get("users", [])),
            system_prompt=str(cfg.get("system_prompt", cfg.get("systemPrompt", ""))),
        )
    return result