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

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass, field
from enum import StrEnum
from pathlib import Path
from typing import Any

from yuxi.utils.logging_config import logger


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


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


@dataclass
class SecurityCheckResult:
    allowed: bool
    reject_reason: str | None = None
    reply: str | None = None


@dataclass
class IMessageSecurityConfig:
    dm_policy: DmPolicy = DmPolicy.PAIRING
    group_policy: GroupPolicy = GroupPolicy.ALLOWLIST
    allow_from: list[str] = field(default_factory=list)
    group_allow_from: list[str] = field(default_factory=list)
    require_mention: bool = False


class IMessageSecurityPolicy:
    def __init__(self, config: dict[str, Any], storage_dir: str | None = None, config_writes: bool = True):
        self._config = config
        self._dm_policy = resolve_dm_policy(config)
        self._group_policy = resolve_group_policy(config)
        self._allow_list = _normalize_entries(config.get("allowFrom", config.get("allow_from", [])))
        self._group_allow_list = _normalize_entries(config.get("groupAllowFrom", config.get("group_allow_from", [])))
        self._require_mention = config.get("requireMention", config.get("require_mention", False))
        self._groups_config: dict[str, dict] = config.get("groups", {})
        self._group_allow_from_fallback = config.get(
            "groupAllowFromFallback", config.get("group_allow_from_fallback", False)
        )
        self._storage_dir = storage_dir or str(Path.home() / ".yuxi" / "imessage")
        self._config_writes = config_writes
        self._account_id = config.get("accountId", config.get("account_id", "default"))
        self._context_visibility_mode = config.get(
            "contextVisibilityMode", config.get("context_visibility_mode", "allowlist")
        )
        self._warned_policies: set[str] = set()
        self._pinned_dm_owner: str | None = None

    def _schedule_save(self) -> None:
        if not self._config_writes:
            return
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(self.save_allowlist())
        except RuntimeError:
            pass

    async def save_allowlist(self) -> None:
        dir_path = Path(self._storage_dir) / self._account_id
        dir_path.mkdir(parents=True, exist_ok=True)
        data = {
            "allow_from": list(self._allow_list),
            "group_allow_from": list(self._group_allow_list),
        }
        file_path = dir_path / "allowlist.json"
        file_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")

    async def load_allowlist(self) -> None:
        file_path = Path(self._storage_dir) / self._account_id / "allowlist.json"
        if not file_path.exists():
            return
        try:
            data = json.loads(file_path.read_text(encoding="utf-8"))
            if isinstance(data, dict):
                loaded_allow = _normalize_entries(data.get("allow_from", []))
                loaded_group = _normalize_entries(data.get("group_allow_from", []))
                for entry in loaded_allow:
                    if entry not in self._allow_list:
                        self._allow_list.append(entry)
                for entry in loaded_group:
                    if entry not in self._group_allow_list:
                        self._group_allow_list.append(entry)
        except (json.JSONDecodeError, OSError) as e:
            logger.warning(f"[iMessage/Security] Failed to load allowlist: {e}")

    def check_dm_access(self, sender_handle: str) -> SecurityCheckResult:
        if self._dm_policy == DmPolicy.OPEN:
            return SecurityCheckResult(allowed=True)
        if self._dm_policy == DmPolicy.DISABLED:
            return SecurityCheckResult(
                allowed=False,
                reject_reason="dm_disabled",
                reply="iMessage DM is disabled for this bot.",
            )
        if self._dm_policy == DmPolicy.ALLOWLIST:
            if _match_entry(sender_handle, self._allow_list):
                return SecurityCheckResult(allowed=True)
            return SecurityCheckResult(
                allowed=False,
                reject_reason="not_in_dm_allowlist",
                reply="You are not authorized to message this bot.",
            )
        if self._dm_policy == DmPolicy.PAIRING:
            return SecurityCheckResult(allowed=True)
        return SecurityCheckResult(allowed=False, reject_reason="unknown_policy")

    def check_group_access(self, chat_guid: str) -> SecurityCheckResult:
        group_override = self._groups_config.get(chat_guid, {})
        if group_override.get("enabled") is False:
            return SecurityCheckResult(
                allowed=False,
                reject_reason="group_disabled_by_override",
                reply="This bot is not enabled in this group chat.",
            )

        if self._group_policy == GroupPolicy.OPEN:
            return SecurityCheckResult(allowed=True)
        if self._group_policy == GroupPolicy.DISABLED:
            return SecurityCheckResult(
                allowed=False,
                reject_reason="group_disabled",
                reply="This bot does not participate in group chats.",
            )
        if self._group_policy == GroupPolicy.ALLOWLIST:
            if _match_entry(chat_guid, self._group_allow_list):
                return SecurityCheckResult(allowed=True)
            if self._group_allow_from_fallback and _match_entry(chat_guid, self._allow_list):
                return SecurityCheckResult(allowed=True)
            return SecurityCheckResult(
                allowed=False,
                reject_reason="group_not_in_allowlist",
                reply="This bot is not authorized in this group chat.",
            )
        return SecurityCheckResult(allowed=False, reject_reason="unknown_policy")

    def check_mention_required(self, chat_guid: str, mentions, bot_handle: str) -> SecurityCheckResult:
        group_override = self._groups_config.get(chat_guid, {})
        require_mention = group_override.get("require_mention", self._require_mention)
        if not require_mention:
            return SecurityCheckResult(allowed=True)

        mentioned_ids: list[str] = []
        if mentions:
            mentioned_ids = mentions.mentioned_user_ids or []
        elif isinstance(mentions, list):
            mentioned_ids = mentions
        elif isinstance(mentions, dict):
            mentioned_ids = mentions.get("mentioned_user_ids", [])

        bot_clean = bot_handle.strip().replace("+", "").replace(" ", "")
        for uid in mentioned_ids:
            uid_clean = uid.strip().replace("+", "").replace(" ", "")
            if uid_clean == bot_clean:
                return SecurityCheckResult(allowed=True)

        return SecurityCheckResult(
            allowed=False,
            reject_reason="mention_required",
            reply=None,
        )

    def collect_warnings(self) -> list[str]:
        warnings: list[str] = []
        if self._dm_policy == DmPolicy.OPEN:
            warnings.append(
                "dmPolicy='open' — anyone can message this bot; consider using 'pairing' or 'allowlist' for security"
            )
        if self._dm_policy == DmPolicy.ALLOWLIST and not self._allow_list:
            warnings.append("dmPolicy='allowlist' but allowFrom is empty — no one can message the bot")
        if self._group_policy == GroupPolicy.OPEN and not self._group_allow_list:
            warnings.append("groupPolicy='open' but no groupAllowFrom configured — consider adding a group allowlist")
        if self._group_policy == GroupPolicy.ALLOWLIST and not self._group_allow_list:
            warnings.append("groupPolicy='allowlist' but groupAllowFrom is empty — no group can interact")
        if self._group_allow_from_fallback:
            warnings.append(
                "groupAllowFromFallback=true — group allowlist falls back to DM allowlist; "
                "consider disabling this for stricter group access control"
            )
        return warnings

    def add_to_allow_list(self, handle: str) -> None:
        handle_clean = _normalize_handle(handle)
        if handle_clean not in self._allow_list:
            self._allow_list.append(handle_clean)
            self._schedule_save()

    def remove_from_allow_list(self, handle: str) -> bool:
        handle_clean = _normalize_handle(handle)
        if handle_clean in self._allow_list:
            self._allow_list.remove(handle_clean)
            self._schedule_save()
            return True
        return False

    def add_to_group_allow_list(self, chat_guid: str) -> None:
        if chat_guid not in self._group_allow_list:
            self._group_allow_list.append(chat_guid)
            self._schedule_save()

    def remove_from_group_allow_list(self, chat_guid: str) -> bool:
        if chat_guid in self._group_allow_list:
            self._group_allow_list.remove(chat_guid)
            self._schedule_save()
            return True
        return False

    def resolve_runtime_group_policy(self, chat_guid: str) -> GroupPolicy:
        group_override = self._groups_config.get(chat_guid, {})
        if "group_policy" in group_override:
            return GroupPolicy(group_override["group_policy"])

        warning_key = f"group_policy_fallback:{chat_guid}"
        if warning_key not in self._warned_policies:
            logger.warning(
                f"[iMessage/Security] No per-group policy for {chat_guid}, "
                f"falling back to default groupPolicy='{self._group_policy.value}'. "
                f"Consider adding a group override config."
            )
            self._warned_policies.add(warning_key)

        return self._group_policy

    def evaluate_context_visibility(
        self,
        chat_guid: str,
        reply_sender_handle: str,
    ) -> bool:
        if self._context_visibility_mode == "open":
            return True

        if self._context_visibility_mode == "disabled":
            return False

        if self._context_visibility_mode == "allowlist":
            chat_type = _resolve_chat_type_from_guid(chat_guid)
            if chat_type == "group":
                if _match_entry(chat_guid, self._group_allow_list):
                    return True
                if self._group_allow_from_fallback and _match_entry(chat_guid, self._allow_list):
                    return True
                if _match_entry(reply_sender_handle, self._allow_list):
                    return True
                return False
            return _match_entry(reply_sender_handle, self._allow_list)

        return False

    def resolve_pinned_dm_owner(self) -> str | None:
        if self._pinned_dm_owner:
            return self._pinned_dm_owner
        for entry in self._allow_list:
            entry_stripped, prefix = _strip_prefix(entry)
            if prefix in ("imessage", "sms", None) and entry_stripped:
                clean = _normalize_handle(entry_stripped)
                if clean and "@" not in clean:
                    self._pinned_dm_owner = clean
                    return clean
        return None

    def should_update_last_route(self, sender_handle: str) -> bool:
        owner = self.resolve_pinned_dm_owner()
        if owner:
            sender_clean = _normalize_handle(sender_handle)
            return sender_clean == owner
        return True

    @property
    def dm_policy(self) -> DmPolicy:
        return self._dm_policy

    @property
    def group_policy(self) -> GroupPolicy:
        return self._group_policy

    @property
    def allow_list(self) -> list[str]:
        return list(self._allow_list)

    @property
    def group_allow_list(self) -> list[str]:
        return list(self._group_allow_list)

    @property
    def require_mention(self) -> bool:
        return self._require_mention


def resolve_dm_policy(config: dict[str, Any]) -> DmPolicy:
    explicit = config.get("dmPolicy", config.get("dm_policy", ""))
    if explicit:
        try:
            return DmPolicy(explicit)
        except ValueError:
            logger.warning(f"[iMessage] Invalid dmPolicy '{explicit}', falling back to auto-detect")
    allow_from = config.get("allowFrom", config.get("allow_from", []))
    if not allow_from or "*" in allow_from:
        return DmPolicy.OPEN
    return DmPolicy.ALLOWLIST


def resolve_group_policy(config: dict[str, Any]) -> GroupPolicy:
    explicit = config.get("groupPolicy", config.get("group_policy", ""))
    if explicit:
        try:
            return GroupPolicy(explicit)
        except ValueError:
            logger.warning(f"[iMessage] Invalid groupPolicy '{explicit}', using ALLOWLIST")
    return GroupPolicy.ALLOWLIST


def _normalize_handle(handle: str) -> str:
    return handle.strip().replace(" ", "").lstrip("+")


def _normalize_entries(entries: list[str]) -> list[str]:
    return [_normalize_handle(e) for e in entries if e]


PREFIX_MAP: dict[str, str] = {
    "chat_id:": "chat_id",
    "chat_guid:": "chat_guid",
    "chat_identifier:": "chat_identifier",
    "imessage:": "imessage",
    "sms:": "sms",
    "auto:": "auto",
}


def _strip_prefix(entry: str) -> tuple[str, str | None]:
    entry_lower = entry.lower()
    for prefix, name in PREFIX_MAP.items():
        if entry_lower.startswith(prefix):
            return entry[len(prefix) :], name
    return entry, None


def _match_entry(target: str, allow_list: list[str]) -> bool:
    if "*" in allow_list:
        return True
    target_clean = _normalize_handle(target)
    for entry in allow_list:
        entry_stripped, prefix = _strip_prefix(entry)
        if prefix == "chat_id" and target_clean in [_normalize_handle(entry_stripped), entry_stripped]:
            return True
        if prefix in ("chat_guid", "chat_identifier") and entry_stripped == target_clean:
            return True
        if prefix in ("imessage", "sms", "auto", None):
            entry_clean = _normalize_handle(entry_stripped)
            if entry_clean == target_clean:
                return True
    return False


def _resolve_chat_type_from_guid(chat_guid: str) -> str:
    if ";-;" in chat_guid or "@chat" in chat_guid:
        return "group"
    return "direct"