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

"""Security policy for Synology Chat access control.

Controls which users can interact with the bot via DM policy, group policy,
and allowlist mechanisms. Includes rate limiting, security warnings
collection, wildcard support, and pairing mode.
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any

from yuxi.channels.models import ChannelMessage, ChatType
from yuxi.utils.logging_config import logger

_DEFAULT_RATE_LIMIT_PER_MINUTE = 30
_DEFAULT_RATE_WINDOW_SECONDS = 60


@dataclass
class _RateWindow:
    count: int = 0
    window_start: float = 0.0


class RateLimiter:
    def __init__(
        self, max_requests: int = _DEFAULT_RATE_LIMIT_PER_MINUTE, window_seconds: int = _DEFAULT_RATE_WINDOW_SECONDS
    ):
        self._max_requests = max_requests
        self._window_seconds = window_seconds
        self._buckets: dict[str, _RateWindow] = {}
        self._gc_counter = 0
        self._GC_INTERVAL = 1000

    def allow(self, key: str) -> bool:
        now = time.monotonic()
        bucket = self._buckets.get(key)
        if bucket is None or now - bucket.window_start > self._window_seconds:
            self._buckets[key] = _RateWindow(count=1, window_start=now)
            self._maybe_gc(now)
            return True
        if bucket.count < self._max_requests:
            bucket.count += 1
            return True
        return False

    def _maybe_gc(self, now: float) -> None:
        self._gc_counter += 1
        if self._gc_counter >= self._GC_INTERVAL:
            self._gc_counter = 0
            cutoff = now - self._window_seconds * 3
            stale = [k for k, v in self._buckets.items() if v.window_start < cutoff]
            for k in stale:
                del self._buckets[k]

    def remaining(self, key: str) -> int:
        now = time.monotonic()
        bucket = self._buckets.get(key)
        if bucket is None or now - bucket.window_start > self._window_seconds:
            return self._max_requests
        return max(0, self._max_requests - bucket.count)


class SynologyChatSecurityPolicy:
    def __init__(self, config: dict[str, Any], account_id: str = ""):
        sec = config.get("security", {})
        if not isinstance(sec, dict):
            sec = {}
        self._dm_policy = sec.get("dm_policy", "open")
        self._group_policy = sec.get("group_policy", "allowlist")
        self._allow_from: list[str] = sec.get("allow_from", [])
        self._group_allow_from: list[str] = sec.get("group_allow_from", [])
        self._pending_pairing: set[str] = set()

        rate_cfg = config.get("rate_limit", {})
        if not isinstance(rate_cfg, dict):
            rate_cfg = {}
        self._rate_limiter = RateLimiter(
            max_requests=rate_cfg.get("max_per_minute", _DEFAULT_RATE_LIMIT_PER_MINUTE),
            window_seconds=rate_cfg.get("window_seconds", _DEFAULT_RATE_WINDOW_SECONDS),
        )
        self._account_id = account_id

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

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

    def check(self, msg: ChannelMessage) -> bool:
        user_id = msg.identity.channel_user_id
        rate_key = f"{self._account_id}:{user_id}" if self._account_id else user_id
        if not self._rate_limiter.allow(rate_key):
            logger.warning(f"[SynologyChat] Rate limit exceeded for user {user_id}")
            return False
        if msg.chat_type == ChatType.DIRECT:
            return self._check_dm(msg)
        return self._check_group(msg)

    def check_rate(self, user_id: str) -> bool:
        rate_key = f"{self._account_id}:{user_id}" if self._account_id else user_id
        return self._rate_limiter.allow(rate_key)

    def _check_dm(self, msg: ChannelMessage) -> bool:
        user_id = msg.identity.channel_user_id

        if self._dm_policy == "open":
            return True
        if self._dm_policy == "disabled":
            logger.debug(f"[SynologyChat] DM blocked (policy=disabled) for user {user_id}")
            return False
        if self._dm_policy == "allowlist":
            allowed = self._is_allowed(user_id, self._allow_from)
            if not allowed:
                logger.debug(f"[SynologyChat] DM blocked (allowlist) for user {user_id}")
            return allowed
        if self._dm_policy == "pairing":
            if self._is_allowed(user_id, self._allow_from):
                return True
            self._pending_pairing.add(user_id)
            logger.info(f"[SynologyChat] DM pairing required for user {user_id}")
            return False
        logger.warning(f"[SynologyChat] Unknown dm_policy '{self._dm_policy}', falling back to open")
        return True

    def _check_group(self, msg: ChannelMessage) -> bool:
        user_id = msg.identity.channel_user_id
        chat_id = msg.identity.channel_chat_id

        if self._group_policy == "open":
            return True
        if self._group_policy == "disabled":
            logger.debug(f"[SynologyChat] Group message blocked (policy=disabled) from user {user_id}")
            return False
        if self._group_policy == "allowlist":
            if self._is_allowed(user_id, self._group_allow_from):
                return True
            logger.debug(f"[SynologyChat] Group message blocked (allowlist) from user {user_id} in chat {chat_id}")
            return False
        logger.warning(f"[SynologyChat] Unknown group_policy '{self._group_policy}', falling back to open")
        return True

    def _is_allowed(self, user_id: str, allowlist: list[str]) -> bool:
        if "*" in allowlist:
            return True
        return user_id in allowlist

    def is_pairing_required(self, user_id: str) -> bool:
        return self._dm_policy == "pairing" and user_id in self._pending_pairing

    def approve_pairing(self, user_id: str) -> None:
        self._pending_pairing.discard(user_id)
        if user_id not in self._allow_from:
            self._allow_from.append(user_id)

    def deny_pairing(self, user_id: str) -> None:
        self._pending_pairing.discard(user_id)

    def add_allowed_user(self, user_id: str, scope: str = "dm") -> bool:
        if scope == "dm":
            target = self._allow_from
        elif scope == "group":
            target = self._group_allow_from
        else:
            return False
        if user_id not in target:
            target.append(user_id)
            logger.info(f"[SynologyChat] Added user {user_id} to {scope} allow_from")
            return True
        return False

    def collect_security_warnings(self) -> list[str]:
        warnings: list[str] = []
        if self._dm_policy == "open":
            warnings.append("dm_policy is 'open': any user can DM the bot without restriction")
        if self._group_policy == "open":
            warnings.append("group_policy is 'open': any group member can trigger the bot")
        if not self._allow_from and self._dm_policy == "allowlist":
            warnings.append("dm_policy is 'allowlist' but allow_from is empty — no users permitted")
        if not self._group_allow_from and self._group_policy == "allowlist":
            warnings.append("group_policy is 'allowlist' but group_allow_from is empty — no users permitted")
        if self._dm_policy == "pairing" and not self._allow_from:
            warnings.append("pairing mode active but initial allow_from is empty — all users require approval")
        return warnings

    def rebuild_rate_limiter(self, max_per_minute: int | None = None, window_seconds: int | None = None) -> None:
        if max_per_minute is not None or window_seconds is not None:
            new_max = max_per_minute if max_per_minute is not None else self._rate_limiter._max_requests
            new_window = window_seconds if window_seconds is not None else self._rate_limiter._window_seconds
            self._rate_limiter = RateLimiter(max_requests=new_max, window_seconds=new_window)
            logger.info(f"[SynologyChat] Rate limiter rebuilt: max_per_minute={new_max}, window_seconds={new_window}")