from __future__ import annotations import hmac import logging import re import time from collections import OrderedDict from dataclasses import dataclass from yuxi.channel.extensions.synology_chat.types import ( INPUT_MAX_LENGTH, DM_POLICY_ALLOWLIST, DM_POLICY_DISABLED, DM_POLICY_OPEN, ResolvedSynologyChatAccount, ) logger = logging.getLogger(__name__) INVALID_TOKEN_WINDOW_MS = 60_000 INVALID_TOKEN_MAX_TRACKED_KEYS = 5000 INVALID_TOKEN_MAX_REQUESTS = 10 DANGEROUS_PATTERNS: list[tuple[re.Pattern, str]] = [ (re.compile(r"ignore\s+(all\s+)?(previous|prior|above)\s+(instructions?|prompts?)", re.IGNORECASE), "[FILTERED]"), (re.compile(r"you\s+are\s+now\s+", re.IGNORECASE), "[FILTERED]"), (re.compile(r"system:\s*", re.IGNORECASE), "[FILTERED]"), (re.compile(r"<\|.*?\|>"), "[FILTERED]"), ] def validate_token(received: str, expected: str) -> bool: if not received or not expected: return False return hmac.compare_digest(received.encode(), expected.encode()) def authorize_dm(account: ResolvedSynologyChatAccount, user_id: str) -> tuple[bool, str]: policy = account.dm_policy allowed = account.allowed_user_ids if policy == DM_POLICY_DISABLED: return False, "DM is disabled" if policy == DM_POLICY_OPEN: if "*" in allowed: return True, "" if not allowed: return False, "not-allowlisted" return user_id in allowed, "not-allowlisted" if policy == DM_POLICY_ALLOWLIST: if not allowed: return False, "allowlist-empty" return user_id in allowed, "not-allowlisted" return False, "unknown-policy" def sanitize_input(text: str) -> str: for pattern, replacement in DANGEROUS_PATTERNS: text = pattern.sub(replacement, text) if len(text) > INPUT_MAX_LENGTH: text = text[:INPUT_MAX_LENGTH] + " ... [truncated]" return text def collect_security_warnings(account: ResolvedSynologyChatAccount) -> list[str]: warnings: list[str] = [] if not account.token: warnings.append("token is not configured. The webhook will reject all requests.") if not account.incoming_url: warnings.append("incomingUrl is not configured. The bot cannot send replies.") if account.allow_insecure_ssl: warnings.append("SSL verification is disabled. Only use for local NAS with self-signed certificates.") if account.dangerously_allow_name_matching: warnings.append("dangerouslyAllowNameMatching re-enables mutable username/nickname recipient matching.") if account.dm_policy == DM_POLICY_OPEN and not account.allowed_user_ids: warnings.append('Add allowedUserIds=["*"] for public DMs or set explicit user IDs.') if account.dm_policy == DM_POLICY_ALLOWLIST and not account.allowed_user_ids: warnings.append('Add users or set dmPolicy="open" with allowedUserIds=["*"].') return warnings @dataclass class _WindowState: count: int window_start_ms: float class InvalidTokenRateLimiter: def __init__(self, max_requests: int = INVALID_TOKEN_MAX_REQUESTS): self._max_requests = max_requests self._state: OrderedDict[str, _WindowState] = OrderedDict() def allow(self, key: str) -> bool: now = time.monotonic() * 1000 state = self._state.get(key) if state is None or (now - state.window_start_ms) >= INVALID_TOKEN_WINDOW_MS: self._state[key] = _WindowState(count=1, window_start_ms=now) self._touch(key) return True state.count += 1 self._touch(key) return state.count <= self._max_requests def _touch(self, key: str) -> None: self._state.move_to_end(key) while len(self._state) > INVALID_TOKEN_MAX_TRACKED_KEYS: self._state.popitem(last=False) def clear(self) -> None: self._state.clear() class UserRateLimiter: def __init__(self, max_per_window: int = 30, window_ms: int = 60_000): self._max = max_per_window self._window_ms = window_ms self._state: dict[str, _WindowState] = {} def allow(self, user_id: str) -> bool: now = time.monotonic() * 1000 state = self._state.get(user_id) if state is None or (now - state.window_start_ms) >= self._window_ms: self._state[user_id] = _WindowState(count=1, window_start_ms=now) return True state.count += 1 return state.count <= self._max def clear(self) -> None: self._state.clear() def max_requests(self) -> int: return self._max _rate_limiters: dict[str, UserRateLimiter] = {} _invalid_token_limiters: dict[str, InvalidTokenRateLimiter] = {} def get_rate_limiter(account: ResolvedSynologyChatAccount) -> UserRateLimiter: rl = _rate_limiters.get(account.account_id) if rl is None or rl.max_requests() != account.rate_limit_per_minute: if rl: rl.clear() rl = UserRateLimiter(max_per_window=account.rate_limit_per_minute) _rate_limiters[account.account_id] = rl return rl def get_invalid_token_limiter(account_id: str) -> InvalidTokenRateLimiter: limiter = _invalid_token_limiters.get(account_id) if limiter is None: limiter = InvalidTokenRateLimiter() _invalid_token_limiters[account_id] = limiter return limiter