ForcePilot/backend/package/yuxi/channel/extensions/synology_chat/security.py

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