416 lines
17 KiB
Python
416 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from yuxi.channel.constants import InboundRejectionReason, channel_rate_limit_key
|
|
from yuxi.channel.plugins.protocol import InboundMessage, OutboundMessage
|
|
from yuxi.channel.security.models import DmPolicy, GroupPolicy, RateLimitPolicy
|
|
from yuxi.channel.security.pairing import PairingChecker
|
|
from yuxi.channel.security.policy import SecurityCheckResult, SecurityPolicy
|
|
from yuxi.channel.security.registry import SecurityCheckerRegistry, SecurityContext
|
|
|
|
|
|
class _MockAllowlistChecker:
|
|
name = "allowlist"
|
|
default_priority = 100
|
|
|
|
def __init__(self, mock: MagicMock) -> None:
|
|
self._mock = mock
|
|
|
|
async def check(self, ctx: SecurityContext) -> object:
|
|
policy = self._resolve_policy(ctx)
|
|
allowed = self._mock.check(policy, ctx.inbound, resolved_sender_id=ctx.resolved_sender_id)
|
|
if allowed:
|
|
return None
|
|
|
|
reason = (
|
|
InboundRejectionReason.DM_NOT_ALLOWED
|
|
if ctx.chat_type == "private"
|
|
else InboundRejectionReason.GROUP_NOT_ALLOWED
|
|
)
|
|
return SecurityCheckResult(allowed=False, reason=reason)
|
|
|
|
def _resolve_policy(self, ctx: SecurityContext) -> object:
|
|
if ctx.chat_type == "private":
|
|
policy = ctx.plugin.resolve_dm_policy(ctx.config, ctx.inbound.account_id)
|
|
return policy if policy is not None else DmPolicy()
|
|
policy = ctx.plugin.resolve_group_policy(ctx.config, ctx.inbound.account_id)
|
|
return policy if policy is not None else GroupPolicy()
|
|
|
|
|
|
class _MockBotLoopChecker:
|
|
name = "bot_loop"
|
|
default_priority = 300
|
|
|
|
def __init__(self, mock: MagicMock) -> None:
|
|
self._mock = mock
|
|
|
|
async def check(self, ctx: SecurityContext) -> object:
|
|
actor_id = ctx.resolved_sender_id or ctx.inbound.sender_id or ctx.inbound.peer_id or ""
|
|
is_bot = False
|
|
if isinstance(ctx.inbound.raw_event, dict):
|
|
is_bot = bool(ctx.inbound.raw_event.get("is_bot", False))
|
|
if self._mock.check(actor_id, is_bot):
|
|
return None
|
|
return SecurityCheckResult(allowed=False, reason=InboundRejectionReason.BOT_LOOP_DETECTED)
|
|
|
|
|
|
class _MockRateLimitChecker:
|
|
name = "rate_limit"
|
|
default_priority = 400
|
|
|
|
def __init__(self, mock: AsyncMock) -> None:
|
|
self._mock = mock
|
|
|
|
async def check(self, ctx: SecurityContext) -> object:
|
|
rate_policy = ctx.plugin.resolve_rate_limit_policy(ctx.config, ctx.inbound.account_id)
|
|
max_requests = (
|
|
rate_policy.max_requests_per_minute
|
|
if rate_policy
|
|
else ctx.config.get("rate_limit", {}).get("max_requests_per_minute", 60)
|
|
)
|
|
actor_id = ctx.resolved_sender_id or ctx.inbound.sender_id or ctx.inbound.peer_id or ""
|
|
rate_key = channel_rate_limit_key(ctx.inbound.channel_type, ctx.inbound.account_id or "", actor_id)
|
|
if await self._mock.is_allowed(rate_key, max_requests, window_seconds=60):
|
|
return None
|
|
return SecurityCheckResult(allowed=False, reason=InboundRejectionReason.RATE_LIMITED)
|
|
|
|
|
|
@pytest.fixture
|
|
def plugin() -> MagicMock:
|
|
p = MagicMock()
|
|
p.resolve_dm_policy.return_value = DmPolicy(mode="open")
|
|
p.resolve_group_policy.return_value = None
|
|
p.resolve_rate_limit_policy.return_value = None
|
|
return p
|
|
|
|
|
|
@pytest.fixture
|
|
def deps() -> dict[str, MagicMock | AsyncMock]:
|
|
identity = MagicMock()
|
|
identity.resolve.return_value = None
|
|
return {
|
|
"allowlist": MagicMock(),
|
|
"pairing": AsyncMock(),
|
|
"bot_loop": MagicMock(),
|
|
"rate_limit": AsyncMock(),
|
|
"identity": identity,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def policy(deps: dict[str, MagicMock | AsyncMock]) -> SecurityPolicy:
|
|
registry = SecurityCheckerRegistry()
|
|
registry.register(_MockAllowlistChecker(deps["allowlist"]))
|
|
registry.register(PairingChecker(manager=deps["pairing"]))
|
|
registry.register(_MockBotLoopChecker(deps["bot_loop"]))
|
|
registry.register(_MockRateLimitChecker(deps["rate_limit"]))
|
|
return SecurityPolicy(registry, identity_link_resolver=deps["identity"])
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_logger(monkeypatch) -> None:
|
|
monkeypatch.setattr("yuxi.channel.security.pairing.logger", MagicMock())
|
|
|
|
|
|
def make_inbound(
|
|
*,
|
|
chat_type: str = "private",
|
|
sender_id: str | None = "u1",
|
|
peer_id: str | None = "p1",
|
|
account_id: str = "a1",
|
|
channel_type: str = "feishu",
|
|
raw_event: dict | None = None,
|
|
is_at_bot: bool = False,
|
|
session_key: str | None = None,
|
|
) -> InboundMessage:
|
|
return InboundMessage(
|
|
channel_type=channel_type,
|
|
account_id=account_id,
|
|
sender_id=sender_id,
|
|
peer_id=peer_id,
|
|
chat_type=chat_type,
|
|
is_at_bot=is_at_bot,
|
|
raw_event=raw_event or {},
|
|
session_key=session_key,
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSecurityPolicyDm:
|
|
async def test_dm_allowed_passes_all_checks(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
inbound = make_inbound()
|
|
deps["identity"].resolve.return_value = "canonical_u1"
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = True
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is True
|
|
assert result.reason is None
|
|
assert result.pairing_code is None
|
|
deps["identity"].resolve.assert_called_once_with("feishu", "u1")
|
|
deps["allowlist"].check.assert_called_once_with(
|
|
DmPolicy(mode="open"), inbound, resolved_sender_id="canonical_u1"
|
|
)
|
|
deps["bot_loop"].check.assert_called_once_with("canonical_u1", False)
|
|
deps["rate_limit"].is_allowed.assert_awaited_once_with(
|
|
channel_rate_limit_key("feishu", "a1", "canonical_u1"), 60, window_seconds=60
|
|
)
|
|
|
|
async def test_dm_not_allowed_triggers_pairing_code(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
plugin.resolve_dm_policy.return_value = DmPolicy(mode="allow_from", allow_list=[])
|
|
deps["allowlist"].check.return_value = False
|
|
deps["pairing"].generate_pairing_code.return_value = "654321"
|
|
inbound = make_inbound(peer_id="p1")
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is False
|
|
assert result.reason == InboundRejectionReason.DM_PAIRING_REQUIRED
|
|
assert result.pairing_code == "654321"
|
|
deps["pairing"].generate_pairing_code.assert_awaited_once_with("feishu", "a1", "p1", pairing_mode="code")
|
|
|
|
async def test_dm_not_allowed_without_build_pairing_qr_reply_uses_default_qr_reply(
|
|
self, policy: SecurityPolicy, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
"""插件未实现 build_pairing_qr_reply 时，应返回默认 OutboundMessage 降级回复。"""
|
|
plugin = MagicMock(spec=["resolve_dm_policy", "resolve_group_policy", "resolve_rate_limit_policy"])
|
|
plugin.resolve_dm_policy.return_value = DmPolicy(mode="allow_from", allow_list=[])
|
|
plugin.resolve_group_policy.return_value = None
|
|
plugin.resolve_rate_limit_policy.return_value = None
|
|
assert not hasattr(plugin, "build_pairing_qr_reply")
|
|
|
|
deps["allowlist"].check.return_value = False
|
|
deps["pairing"].generate_pairing_code.return_value = "654321"
|
|
config = {"pairing": {"mode": "qr"}}
|
|
inbound = make_inbound(peer_id="p1")
|
|
|
|
result = await policy.check(config, plugin, inbound)
|
|
|
|
assert result.allowed is False
|
|
assert result.reason == InboundRejectionReason.DM_PAIRING_REQUIRED
|
|
assert result.pairing_code == "654321"
|
|
assert result.qr_content == "654321"
|
|
assert isinstance(result.qr_reply, OutboundMessage)
|
|
assert result.qr_reply.content == "654321"
|
|
assert result.qr_reply.content_type == "text"
|
|
|
|
async def test_dm_deny_mode_does_not_generate_pairing_code(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
plugin.resolve_dm_policy.return_value = DmPolicy(mode="deny")
|
|
deps["allowlist"].check.return_value = False
|
|
inbound = make_inbound()
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is False
|
|
assert result.reason == InboundRejectionReason.DM_NOT_ALLOWED
|
|
assert result.pairing_code is None
|
|
deps["pairing"].generate_pairing_code.assert_not_awaited()
|
|
|
|
async def test_dm_pairing_failure_falls_back_to_dm_not_allowed(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
plugin.resolve_dm_policy.return_value = DmPolicy(mode="allow_from", allow_list=[])
|
|
deps["allowlist"].check.return_value = False
|
|
deps["pairing"].generate_pairing_code.side_effect = RuntimeError("db down")
|
|
inbound = make_inbound()
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is False
|
|
assert result.reason == InboundRejectionReason.DM_NOT_ALLOWED
|
|
assert result.pairing_code is None
|
|
|
|
async def test_dm_defaults_to_open_policy_when_plugin_returns_none(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
plugin.resolve_dm_policy.return_value = None
|
|
deps["identity"].resolve.return_value = None
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = True
|
|
inbound = make_inbound()
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is True
|
|
deps["allowlist"].check.assert_called_once_with(DmPolicy(), inbound, resolved_sender_id=None)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSecurityPolicyGroup:
|
|
async def test_group_not_allowed(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
inbound = make_inbound(chat_type="group", session_key="g1")
|
|
plugin.resolve_group_policy.return_value = GroupPolicy(deny_groups=["g1"])
|
|
deps["allowlist"].check.return_value = False
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is False
|
|
assert result.reason == InboundRejectionReason.GROUP_NOT_ALLOWED
|
|
plugin.resolve_group_policy.assert_called_once_with({}, "a1")
|
|
|
|
async def test_group_defaults_policy_when_plugin_returns_none(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
inbound = make_inbound(chat_type="group", session_key="g1")
|
|
deps["identity"].resolve.return_value = "u1"
|
|
plugin.resolve_group_policy.return_value = None
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = True
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is True
|
|
deps["allowlist"].check.assert_called_once_with(GroupPolicy(), inbound, resolved_sender_id="u1")
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSecurityPolicyBotLoopAndRateLimit:
|
|
async def test_bot_loop_detected(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
deps["identity"].resolve.return_value = "canonical_u1"
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = False
|
|
inbound = make_inbound(raw_event={"is_bot": True})
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is False
|
|
assert result.reason == InboundRejectionReason.BOT_LOOP_DETECTED
|
|
deps["bot_loop"].check.assert_called_once_with("canonical_u1", True)
|
|
|
|
async def test_raw_event_not_dict_treated_as_not_bot(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = True
|
|
inbound = make_inbound(raw_event="not_a_dict")
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is True
|
|
deps["bot_loop"].check.assert_called_once_with("u1", False)
|
|
|
|
async def test_rate_limited(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = False
|
|
inbound = make_inbound()
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is False
|
|
assert result.reason == InboundRejectionReason.RATE_LIMITED
|
|
|
|
async def test_rate_limit_uses_config_when_plugin_returns_none(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = True
|
|
config = {"rate_limit": {"max_requests_per_minute": 42}}
|
|
inbound = make_inbound()
|
|
|
|
result = await policy.check(config, plugin, inbound)
|
|
|
|
assert result.allowed is True
|
|
deps["rate_limit"].is_allowed.assert_awaited_once_with(
|
|
channel_rate_limit_key("feishu", "a1", "u1"), 42, window_seconds=60
|
|
)
|
|
|
|
async def test_rate_limit_defaults_to_60(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = True
|
|
inbound = make_inbound()
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is True
|
|
deps["rate_limit"].is_allowed.assert_awaited_once_with(
|
|
channel_rate_limit_key("feishu", "a1", "u1"), 60, window_seconds=60
|
|
)
|
|
|
|
async def test_rate_limit_uses_plugin_policy(
|
|
self, policy: SecurityPolicy, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
plugin.resolve_rate_limit_policy.return_value = RateLimitPolicy(max_requests_per_minute=99)
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = True
|
|
inbound = make_inbound()
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is True
|
|
deps["rate_limit"].is_allowed.assert_awaited_once_with(
|
|
channel_rate_limit_key("feishu", "a1", "u1"), 99, window_seconds=60
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSecurityPolicyIdentityFallback:
|
|
async def test_no_identity_resolver_uses_original_sender_id(
|
|
self, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
registry = SecurityCheckerRegistry()
|
|
registry.register(_MockAllowlistChecker(deps["allowlist"]))
|
|
registry.register(PairingChecker(manager=deps["pairing"]))
|
|
registry.register(_MockBotLoopChecker(deps["bot_loop"]))
|
|
registry.register(_MockRateLimitChecker(deps["rate_limit"]))
|
|
policy = SecurityPolicy(registry, identity_link_resolver=None)
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = True
|
|
inbound = make_inbound()
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is True
|
|
deps["allowlist"].check.assert_called_once_with(DmPolicy(mode="open"), inbound, resolved_sender_id=None)
|
|
deps["bot_loop"].check.assert_called_once_with("u1", False)
|
|
deps["rate_limit"].is_allowed.assert_awaited_once_with(
|
|
channel_rate_limit_key("feishu", "a1", "u1"), 60, window_seconds=60
|
|
)
|
|
|
|
async def test_actor_id_falls_back_to_peer_id(
|
|
self, plugin: MagicMock, deps: dict[str, MagicMock | AsyncMock]
|
|
) -> None:
|
|
registry = SecurityCheckerRegistry()
|
|
registry.register(_MockAllowlistChecker(deps["allowlist"]))
|
|
registry.register(PairingChecker(manager=deps["pairing"]))
|
|
registry.register(_MockBotLoopChecker(deps["bot_loop"]))
|
|
registry.register(_MockRateLimitChecker(deps["rate_limit"]))
|
|
policy = SecurityPolicy(registry, identity_link_resolver=None)
|
|
deps["allowlist"].check.return_value = True
|
|
deps["bot_loop"].check.return_value = True
|
|
deps["rate_limit"].is_allowed.return_value = True
|
|
inbound = make_inbound(sender_id=None, peer_id="peer_1")
|
|
|
|
result = await policy.check({}, plugin, inbound)
|
|
|
|
assert result.allowed is True
|
|
deps["bot_loop"].check.assert_called_once_with("peer_1", False)
|
|
deps["rate_limit"].is_allowed.assert_awaited_once_with(
|
|
channel_rate_limit_key("feishu", "a1", "peer_1"), 60, window_seconds=60
|
|
)
|