ForcePilot/backend/package/yuxi/channels/policy/security_policy.py

from __future__ import annotations

import fnmatch
from dataclasses import dataclass, field
from enum import StrEnum

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"


class RejectReason(StrEnum):
    DM_DISABLED = "dm_policy_disabled"
    DM_OPEN_PASS = "dm_policy_open_pass"
    DM_ALLOWLISTED = "dm_policy_allowlisted"
    DM_NOT_ALLOWLISTED = "dm_policy_not_allowlisted"
    DM_EMPTY_ALLOWLIST = "dm_policy_empty_allowlist"
    DM_PAIRING_REQUIRED = "dm_pairing_required"
    DM_PAIRING_RATE_LIMITED = "dm_pairing_rate_limited"
    DM_PAIRING_FULL = "dm_pairing_full"
    DM_PAIRING_PENDING = "dm_pairing_pending"
    GROUP_DISABLED = "group_policy_disabled"
    GROUP_OPEN_PASS = "group_policy_open_pass"
    GROUP_ALLOWLISTED = "group_policy_allowlisted"
    GROUP_NOT_ALLOWLISTED = "group_policy_not_allowlisted"
    GROUP_EMPTY_ALLOWLIST = "group_policy_empty_allowlist"
    UNKNOWN_SENDER = "unknown_sender"
    BLOCKED_SENDER = "blocked_sender"


@dataclass
class WildcardAllowlistMatcher:
    entries: list[str] = field(default_factory=list)

    def __post_init__(self):
        self._exact: set[str] = set()
        self._patterns: list[str] = []
        self._has_wildcard: bool = False
        for entry in self.entries:
            entry = entry.strip()
            if not entry:
                continue
            if entry == "*":
                self._has_wildcard = True
            elif "*" in entry or "?" in entry:
                self._patterns.append(entry)
            else:
                self._exact.add(entry)
                stripped = entry.removeprefix("+")
                if stripped != entry:
                    self._exact.add(stripped)

    def match(self, target: str) -> bool:
        if self._has_wildcard:
            return True
        if target in self._exact:
            return True
        stripped = target.strip().removeprefix("+")
        if stripped in self._exact:
            return True
        for pattern in self._patterns:
            if fnmatch.fnmatch(target, pattern) or fnmatch.fnmatch(stripped, pattern):
                return True
        return False

    @property
    def is_empty(self) -> bool:
        return not self._exact and not self._patterns and not self._has_wildcard

    @classmethod
    def from_config(cls, allow_from: list[str]) -> WildcardAllowlistMatcher:
        return cls(entries=list(allow_from))


@dataclass
class AccessResult:
    allowed: bool
    reason: RejectReason | None = None
    detail: str = ""


class BaseSecurityPolicy:
    def __init__(self, config: dict):
        self._config = config
        self._dm_policy = self._resolve_dm_policy(config)
        self._group_policy = self._resolve_group_policy(config)
        self._dm_matcher = WildcardAllowlistMatcher.from_config(config.get("allowFrom", config.get("allow_from", [])))
        self._group_matcher = WildcardAllowlistMatcher.from_config(
            config.get("groupAllowFrom", config.get("group_allow_from", []))
        )

    def _resolve_sender_id(self, event_data: dict) -> str:
        raise NotImplementedError(f"{self.__class__.__name__} must implement _resolve_sender_id()")

    def _resolve_group_id(self, event_data: dict) -> str:
        raise NotImplementedError(f"{self.__class__.__name__} must implement _resolve_group_id()")

    def _resolve_dm_policy(self, config: dict) -> DmPolicy:
        explicit = config.get("dmPolicy", config.get("dm_policy", ""))
        if explicit:
            try:
                return DmPolicy(explicit)
            except ValueError:
                logger.warning(
                    f"Invalid dmPolicy '{explicit}' in {self.__class__.__name__}, 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(self, config: dict) -> GroupPolicy:
        explicit = config.get("groupPolicy", config.get("group_policy", ""))
        if explicit:
            try:
                return GroupPolicy(explicit)
            except ValueError:
                logger.warning(f"Invalid groupPolicy '{explicit}', using ALLOWLIST")
        return GroupPolicy.ALLOWLIST

    def check_dm_access(self, sender_id: str) -> AccessResult:
        if self._dm_policy == DmPolicy.DISABLED:
            return AccessResult(False, RejectReason.DM_DISABLED)
        if self._dm_policy == DmPolicy.OPEN:
            return AccessResult(True, RejectReason.DM_OPEN_PASS)
        if self._dm_policy == DmPolicy.ALLOWLIST:
            if self._dm_matcher.is_empty:
                return AccessResult(
                    False, RejectReason.DM_EMPTY_ALLOWLIST, "allowFrom is empty — no one can DM the bot"
                )
            if self._dm_matcher.match(sender_id):
                return AccessResult(True, RejectReason.DM_ALLOWLISTED)
            return AccessResult(False, RejectReason.DM_NOT_ALLOWLISTED, f"sender '{sender_id}' not in allowFrom")
        return AccessResult(True, None)

    def check_group_access(self, group_id: str) -> AccessResult:
        if self._group_policy == GroupPolicy.DISABLED:
            return AccessResult(False, RejectReason.GROUP_DISABLED)
        if self._group_policy == GroupPolicy.OPEN:
            return AccessResult(True, RejectReason.GROUP_OPEN_PASS)
        if self._group_policy == GroupPolicy.ALLOWLIST:
            if self._group_matcher.is_empty:
                return AccessResult(
                    False, RejectReason.GROUP_EMPTY_ALLOWLIST, "groupAllowFrom is empty — no groups can interact"
                )
            if self._group_matcher.match(group_id):
                return AccessResult(True, RejectReason.GROUP_ALLOWLISTED)
            return AccessResult(False, RejectReason.GROUP_NOT_ALLOWLISTED, f"group '{group_id}' not in groupAllowFrom")
        return AccessResult(False, RejectReason.GROUP_DISABLED)

    def collect_warnings(self) -> list[str]:
        warnings: list[str] = []
        if self._group_policy == GroupPolicy.OPEN and self._group_matcher.is_empty:
            warnings.append(
                "groupPolicy is 'open' but no groupAllowFrom configured — "
                "any group can interact with the bot. Consider adding an allowlist."
            )
        if self._dm_policy == DmPolicy.OPEN and ("*" not in self._dm_matcher.entries):
            warnings.append(
                "dmPolicy is 'open' but allowFrom is not '*' — "
                "only explicitly listed senders can DM. "
                "To open to everyone, add '*' to allowFrom."
            )
        if self._dm_policy == DmPolicy.ALLOWLIST and self._dm_matcher.is_empty:
            warnings.append(
                "dmPolicy is 'allowlist' but allowFrom is empty — "
                "no one can DM the bot. Add at least one entry to allowFrom."
            )
        return warnings

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

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

    @property
    def dm_allowlist(self) -> WildcardAllowlistMatcher:
        return self._dm_matcher

    @property
    def group_allowlist(self) -> WildcardAllowlistMatcher:
        return self._group_matcher