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

from __future__ import annotations

import logging
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

DM_POLICIES = ("pairing", "allowlist", "open", "disabled")
GROUP_POLICIES = ("open", "allowlist", "disabled")
REPLY_STYLES = ("thread", "top-level")


@dataclass
class MSTeamsDmPolicy:
    mode: str = "pairing"
    allow_from: list[str] = field(default_factory=list)


@dataclass
class MSTeamsGroupPolicy:
    mode: str = "allowlist"
    allow_from: list[str] = field(default_factory=list)
    require_mention: bool = True


@dataclass
class MSTeamsChannelPolicy:
    require_mention: bool = True
    reply_style: str = "thread"
    tools: dict | None = None
    tools_by_sender: list[dict] | None = None


class MSTeamsSecurity:
    def __init__(self, config: dict | None = None):
        self._config = config or {}

    def _channel_cfg(self) -> dict:
        return self._config.get("channels", {}).get("msteams", {}) or {}

    def resolve_dm_policy(self, account: dict | None = None) -> MSTeamsDmPolicy:
        channel_cfg = self._channel_cfg()
        mode = channel_cfg.get("dmPolicy", "pairing")
        allow_from = channel_cfg.get("allowFrom", [])
        if account:
            mode = account.get("dm_policy", mode)
            allow_from = account.get("allow_from", allow_from)
        return MSTeamsDmPolicy(mode=mode, allow_from=allow_from)

    def resolve_group_policy(
        self,
        account: dict | None = None,
        team_id: str | None = None,
        channel_id: str | None = None,
    ) -> MSTeamsGroupPolicy:
        channel_cfg = self._channel_cfg()
        default_mode = channel_cfg.get("groupPolicy", "allowlist")
        default_mention = channel_cfg.get("requireMention", True)
        allow_from = channel_cfg.get("groupAllowFrom", [])

        teams = channel_cfg.get("teams", {}) or {}

        if team_id:
            team_cfg = teams.get(team_id, {}) or {}
            if channel_id and team_cfg:
                channels = team_cfg.get("channels", {}) or {}
                ch_cfg = channels.get(channel_id, {}) or {}
                if ch_cfg:
                    return MSTeamsGroupPolicy(
                        mode=ch_cfg.get("policy", default_mode),
                        require_mention=ch_cfg.get("requireMention", team_cfg.get("requireMention", default_mention)),
                    )
            return MSTeamsGroupPolicy(
                mode=team_cfg.get("policy", default_mode),
                require_mention=team_cfg.get("requireMention", default_mention),
            )

        wildcard = teams.get("*", {}) or {}
        return MSTeamsGroupPolicy(
            mode=wildcard.get("policy", default_mode),
            allow_from=allow_from,
            require_mention=wildcard.get("requireMention", default_mention),
        )

    def is_allowed_dm(self, policy: MSTeamsDmPolicy, peer_id: str) -> tuple[bool, str | None]:
        if policy.mode == "disabled":
            return False, "dm-disabled"
        if policy.mode == "open":
            return True, None
        if policy.mode == "allowlist":
            if _check_allowlist(policy.allow_from, peer_id):
                return True, None
            return False, "not-in-allowlist"
        if policy.mode == "pairing":
            if _check_allowlist(policy.allow_from, peer_id):
                return True, "paired"
            return True, "pairing-required"
        return False, "unknown-policy"

    def is_allowed_group(
        self,
        policy: MSTeamsGroupPolicy,
        peer_id: str,
        is_mentioned: bool = False,
    ) -> tuple[bool, str | None]:
        if policy.mode == "disabled":
            return False, "group-disabled"
        if policy.require_mention and not is_mentioned:
            return False, "mention-required"
        if policy.mode == "open":
            return True, None
        if policy.mode == "allowlist":
            if _check_allowlist(policy.allow_from, peer_id) or not policy.allow_from:
                return True, None
            return False, "not-in-group-allowlist"
        return False, "unknown-policy"

    def resolve_reply_style(
        self,
        team_id: str | None = None,
        channel_id: str | None = None,
    ) -> str:
        channel_cfg = self._channel_cfg()
        default_style = channel_cfg.get("replyStyle", "thread")

        if not team_id:
            return default_style

        teams = channel_cfg.get("teams", {}) or {}
        team_cfg = teams.get(team_id, {}) or {}

        if channel_id:
            channels = team_cfg.get("channels", {}) or {}
            ch_cfg = channels.get(channel_id, {}) or {}
            if "replyStyle" in ch_cfg:
                return ch_cfg["replyStyle"]

        return team_cfg.get("replyStyle", default_style)

    def resolve_tools_policy(
        self,
        team_id: str | None = None,
        channel_id: str | None = None,
        sender_peer_id: str | None = None,
    ) -> dict:
        channel_cfg = self._channel_cfg()
        default_tools = channel_cfg.get("tools", {}) or {}

        if not team_id:
            return default_tools

        teams = channel_cfg.get("teams", {}) or {}
        team_cfg = teams.get(team_id, {}) or {}
        team_tools = team_cfg.get("tools", {}) or {}

        if channel_id:
            channels = team_cfg.get("channels", {}) or {}
            ch_cfg = channels.get(channel_id, {}) or {}
            ch_tools = ch_cfg.get("tools", {}) or {}
            effective = {**default_tools, **team_tools, **ch_tools}
        else:
            effective = {**default_tools, **team_tools}

        tools_by_sender = effective.get("toolsBySender", [])
        if sender_peer_id and tools_by_sender:
            for rule in tools_by_sender:
                if rule.get("sender") == sender_peer_id or rule.get("sender") == "*":
                    return rule

        return effective


def _normalize_peer(peer_id: str) -> str:
    for prefix in ("msteams:", "teams:"):
        if peer_id.startswith(prefix):
            return peer_id[len(prefix) :]
    return str(peer_id)


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