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

import asyncio
import logging
import time
from collections.abc import Awaitable, Callable
from enum import StrEnum

from yuxi.channels.models import ChannelMessage

logger = logging.getLogger(__name__)


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


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


class ReactionNotificationPolicy(StrEnum):
    OFF = "off"
    OWN = "own"
    ALLOWLIST = "allowlist"
    ALL = "all"


class SignalSecurityPolicy:
    def __init__(
        self,
        dm_policy: str = "pairing",
        group_policy: str = "allowlist",
        allow_from: list[str] | None = None,
        group_allow_from: list[str] | None = None,
        require_mention: bool = False,
        reaction_notifications: str = "all",
        reaction_allowlist: list[str] | None = None,
        command_double_auth: bool = True,
    ):
        self.dm_policy = DmPolicy(dm_policy)
        self.group_policy = GroupPolicy(group_policy)
        self._dm_allowlist: set[str] = set(self._expand_allowlist(allow_from or []))
        self._group_allowlist: set[str] = set(group_allow_from or [])
        self._pairing_pending: set[str] = set()
        self.require_mention = require_mention
        self.reaction_notifications = ReactionNotificationPolicy(reaction_notifications)
        self._reaction_allowlist: set[str] = set(reaction_allowlist or [])
        self._pairing_challenge_pending: set[str] = set()
        self._command_double_auth = command_double_auth
        self._store_write_fn: Callable[[str, dict], Awaitable[None]] | None = None
        self._store_read_fn: Callable[[str], Awaitable[dict | None]] | None = None

    @staticmethod
    def _expand_allowlist(entries: list[str]) -> list[str]:
        result = []
        for entry in entries:
            stripped = entry.strip()
            if stripped.startswith("signal:"):
                stripped = stripped.removeprefix("signal:")
            result.append(stripped)
        return result

    def _is_wildcard_match(self, user_id: str) -> bool:
        return "*" in self._dm_allowlist

    def check_dm_permission(self, message: ChannelMessage) -> bool:
        user_id = message.identity.channel_user_id

        match self.dm_policy:
            case DmPolicy.DISABLED:
                logger.info(f"[Signal Security] DM denied (policy=disabled, user={user_id})")
                return False
            case DmPolicy.OPEN:
                return True
            case DmPolicy.ALLOWLIST:
                allowed = user_id in self._dm_allowlist or self._is_wildcard_match(user_id)
                if not allowed:
                    logger.info(f"[Signal Security] DM denied (policy=allowlist, user={user_id})")
                return allowed
            case DmPolicy.PAIRING:
                if user_id in self._dm_allowlist or self._is_wildcard_match(user_id):
                    return True
                self._pairing_pending.add(user_id)
                logger.info(f"[Signal Security] DM denied (policy=pairing, user={user_id}, pending)")
                return False
            case _:
                logger.info(f"[Signal Security] DM denied (policy=unknown, user={user_id})")
                return False

    def check_group_permission(self, message: ChannelMessage) -> bool:
        group_id = message.identity.channel_chat_id

        match self.group_policy:
            case GroupPolicy.DISABLED:
                logger.info(f"[Signal Security] Group denied (policy=disabled, group={group_id})")
                return False
            case GroupPolicy.OPEN:
                return True
            case GroupPolicy.ALLOWLIST:
                allowed = group_id in self._group_allowlist
                if not allowed:
                    logger.info(f"[Signal Security] Group denied (policy=allowlist, group={group_id})")
                return allowed
            case _:
                logger.info(f"[Signal Security] Group denied (policy=unknown, group={group_id})")
                return False

    def check_require_mention(self, message: ChannelMessage) -> bool:
        if not self.require_mention:
            return True
        if message.chat_type.value != "group":
            return True
        if message.mentions and message.mentions.is_bot_mentioned:
            return True
        return False

    def check_reaction_notification(self, user_id: str) -> bool:
        match self.reaction_notifications:
            case ReactionNotificationPolicy.OFF:
                return False
            case ReactionNotificationPolicy.OWN:
                return False
            case ReactionNotificationPolicy.ALLOWLIST:
                return user_id in self._reaction_allowlist
            case ReactionNotificationPolicy.ALL:
                return True
            case _:
                return True

    def record_pairing_challenge(self, user_id: str) -> None:
        self._pairing_challenge_pending.add(user_id)

    def has_pairing_challenge(self, user_id: str) -> bool:
        return user_id in self._pairing_challenge_pending

    def reject_pairing(self, user_id: str) -> None:
        self._pairing_pending.discard(user_id)

    @property
    def pending_pairings(self) -> set[str]:
        return self._pairing_pending.copy()

    def add_to_allowlist(self, target_id: str, target_type: str = "dm") -> None:
        if target_type == "dm":
            self._dm_allowlist.add(target_id)
        elif target_type == "group":
            self._group_allowlist.add(target_id)
        else:
            raise ValueError(f"Unknown target_type: {target_type}, expected 'dm' or 'group'")

        if self._store_write_fn:
            try:
                asyncio.ensure_future(
                    self._store_write_fn(
                        f"signal:pairing:{target_type}:{target_id}",
                        {
                            "user_id": target_id,
                            "target_type": target_type,
                            "added_at": int(time.time()),
                        },
                    )
                )
            except RuntimeError:
                logger.debug("No running event loop, skip persist add_to_allowlist")

    def remove_from_allowlist(self, target_id: str, target_type: str = "dm") -> None:
        if target_type == "dm":
            self._dm_allowlist.discard(target_id)
        elif target_type == "group":
            self._group_allowlist.discard(target_id)
        else:
            raise ValueError(f"Unknown target_type: {target_type}, expected 'dm' or 'group'")

    def check_command_double_auth(self, message: ChannelMessage) -> bool:
        if not self._command_double_auth:
            return True
        user_id = message.identity.channel_user_id
        group_id = message.identity.channel_chat_id

        dm_allowed = user_id in self._dm_allowlist or self._is_wildcard_match(user_id)
        group_allowed = group_id in self._group_allowlist

        if dm_allowed or group_allowed:
            return True

        logger.info(
            f"[Signal Security] Command double-auth denied: user={user_id}, "
            f"group={group_id}, dm_allowed={dm_allowed}, group_allowed={group_allowed}"
        )
        return False

    def set_store_handlers(
        self,
        write_fn: Callable[[str, dict], Awaitable[None]],
        read_fn: Callable[[str], Awaitable[dict | None]],
    ) -> None:
        self._store_write_fn = write_fn
        self._store_read_fn = read_fn

    async def approve_pairing(self, user_id: str) -> None:
        self._dm_allowlist.add(user_id)
        self._pairing_pending.discard(user_id)
        if self._store_write_fn:
            try:
                await self._store_write_fn(
                    f"signal:pairing:dm:{user_id}",
                    {"user_id": user_id, "approved_at": __import__("time").time()},
                )
            except Exception:
                logger.exception("Failed to persist pairing approval to store")
        else:
            logger.debug(f"Store handlers not configured, pairing approval for {user_id} will not survive restart")

    async def load_pairing_store(self) -> None:
        if not self._store_read_fn:
            return
        try:
            data = await self._store_read_fn("signal:pairing:dm:*")
            if data and isinstance(data, dict):
                for key, value in data.items():
                    user_id = value.get("user_id") if isinstance(value, dict) else str(value)
                    if user_id:
                        self._dm_allowlist.add(user_id)
        except Exception:
            logger.exception("Failed to load pairing store")