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

from __future__ import annotations

from dataclasses import dataclass, field
from enum import StrEnum

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 WhatsAppSecurityConfig:
    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)


class WhatsAppSecurityPolicy:
    def __init__(self, config: dict):
        self._config = config
        self._dm_policy = self._resolve_dm_policy()
        self._group_policy = self._resolve_group_policy()
        self._allow_list = self._normalize_allow_list(config.get("allowFrom", config.get("allow_from", [])))
        self._group_allow_list = self._normalize_allow_list(
            config.get("groupAllowFrom", config.get("group_allow_from", []))
        )

    def check_dm_access(self, sender_e164: str) -> tuple[bool, str | None]:
        if self._dm_policy == DmPolicy.DISABLED:
            return False, "dm_disabled"
        if self._dm_policy == DmPolicy.OPEN:
            return True, None
        if self._dm_policy == DmPolicy.ALLOWLIST:
            if self._match_allow_list(sender_e164, self._allow_list):
                return True, None
            return False, "not_in_allowlist"
        if self._dm_policy == DmPolicy.PAIRING:
            return True, None
        return False, "unknown_policy"

    def check_group_access(self, group_jid: str) -> tuple[bool, str | None]:
        if self._group_policy == GroupPolicy.DISABLED:
            return False, "group_disabled"
        if self._group_policy == GroupPolicy.OPEN:
            return True, None
        if self._group_policy == GroupPolicy.ALLOWLIST:
            if self._match_allow_list(group_jid, self._group_allow_list):
                return True, None
            return False, "group_not_in_allowlist"
        return False, "unknown_policy"

    def collect_warnings(self) -> list[str]:
        warnings = []
        if self._group_policy == GroupPolicy.OPEN and not self._group_allow_list:
            warnings.append("groupPolicy is 'open' but no groupAllowFrom configured - consider adding group allowlist")
        if self._dm_policy == DmPolicy.OPEN and ("*" not in self._allow_list):
            warnings.append("dmPolicy is 'open' but allowFrom is not '*' — restricted exposure recommended")
        return warnings

    def apply_config_fixes(self) -> dict:
        fixes = {}
        if self._group_policy == GroupPolicy.OPEN and not self._group_allow_list:
            fixes["group_policy_warning"] = "groupPolicy=open without groupAllowFrom"
        if self._dm_policy == DmPolicy.OPEN and not self._allow_list:
            fixes["dm_policy_warning"] = "dmPolicy=open without allowFrom restriction"
        return fixes

    def add_to_allow_list(self, e164: str) -> None:
        if e164 not in self._allow_list:
            self._allow_list.append(e164)

    def remove_from_allow_list(self, e164: str) -> bool:
        if e164 in self._allow_list:
            self._allow_list.remove(e164)
            return True
        return False

    @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)

    @staticmethod
    def _normalize_allow_list(raw: list[str]) -> list[str]:
        return [entry.strip().removeprefix("+") for entry in raw if entry]

    @staticmethod
    def _match_allow_list(target: str, allow_list: list[str]) -> bool:
        if "*" in allow_list:
            return True
        target_clean = target.removeprefix("+")
        for entry in allow_list:
            if entry == target_clean or entry == target:
                return True
        return False

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

    def _resolve_group_policy(self) -> GroupPolicy:
        explicit = self._config.get("groupPolicy", self._config.get("group_policy", ""))
        if explicit:
            try:
                return GroupPolicy(explicit)
            except ValueError:
                logger.warning(f"Invalid groupPolicy '{explicit}', using ALLOWLIST")
        return GroupPolicy.ALLOWLIST