from __future__ import annotations from datetime import UTC, datetime from types import SimpleNamespace from unittest.mock import AsyncMock import pytest from pyrogram.enums import ChatMemberStatus, MessageEntityType def make_message( *, text: str | None = None, caption: str | None = None, photo: object | None = None, entities: list[object] | None = None, user_id: int = 10, ): return SimpleNamespace( text=text, caption=caption, photo=photo, entities=entities or [], caption_entities=[], chat=SimpleNamespace(id=-100), from_user=SimpleNamespace( id=user_id, is_bot=False, mention=f"user-{user_id}", ), delete=AsyncMock(), ) def grant_bot_permissions(app_modules, *, delete: bool = True, restrict: bool = True): app_modules.app.members[(-100, 999)] = SimpleNamespace( status=ChatMemberStatus.ADMINISTRATOR, privileges=SimpleNamespace( can_delete_messages=delete, can_restrict_members=restrict, ), ) def test_risk_policy_supports_legacy_keywords_and_rejects_conflicting_actions( app_modules, ): service = app_modules.load("wbb.services.blacklist_enforcement") policy = service.normalize_risk_control(None, fallback_keywords=["做爱"]) assert policy == { "enabled": True, "keywords": ["做爱"], "match_images": False, "match_links": False, "actions": ["delete"], "duration_seconds": 3600, } with pytest.raises(service.RiskControlValidationError): service.normalize_risk_control( {"enabled": True, "actions": ["warn", "ban"]}, strict=True, ) def test_risk_match_supports_chinese_images_links_and_normalization(app_modules): service = app_modules.load("wbb.services.blacklist_enforcement") assert service.find_blacklist_match("这句话包含做爱关键词", ["做爱"]) == "做爱" assert service.find_blacklist_match("SPAM message", ["spam"]) == "spam" assert service.find_blacklist_match("普通消息", ["做爱"]) is None policy = service.normalize_risk_control( { "enabled": True, "keywords": [], "match_images": True, "match_links": True, "actions": ["delete"], } ) match = service.detect_risk_match( make_message( text="查看链接", photo=object(), entities=[SimpleNamespace(type=MessageEntityType.URL)], ), policy, ) assert match.trigger_types == ("image", "link") async def test_admin_message_is_deleted_without_member_punishment( app_modules, monkeypatch ): service = app_modules.load("wbb.services.blacklist_enforcement") grant_bot_permissions(app_modules) message = make_message(text="管理员也包含做爱关键词") app_modules.app.members[(-100, 10)] = SimpleNamespace( status=ChatMemberStatus.ADMINISTRATOR ) ban_member = AsyncMock() monkeypatch.setattr(service.app, "ban_chat_member", ban_member, raising=False) result = await service.enforce_risk_message( message, { "enabled": True, "keywords": ["做爱"], "actions": ["delete", "ban"], }, ) message.delete.assert_awaited_once() ban_member.assert_not_awaited() assert result.deleted is True assert result.applied_member_action == "none" async def test_regular_member_photo_can_be_deleted_and_kicked( app_modules, monkeypatch ): service = app_modules.load("wbb.services.blacklist_enforcement") grant_bot_permissions(app_modules) message = make_message(photo=object()) ban_member = AsyncMock() unban_member = AsyncMock() monkeypatch.setattr(service.app, "ban_chat_member", ban_member, raising=False) monkeypatch.setattr(service.app, "unban_chat_member", unban_member, raising=False) result = await service.enforce_risk_message( message, { "enabled": True, "match_images": True, "actions": ["delete", "kick"], }, ) message.delete.assert_awaited_once() ban_member.assert_awaited_once_with(-100, 10) unban_member.assert_awaited_once_with(-100, 10) assert result.applied_member_action == "kick" async def test_delete_failure_does_not_block_selected_ban(app_modules, monkeypatch): service = app_modules.load("wbb.services.blacklist_enforcement") grant_bot_permissions(app_modules) message = make_message(text="做爱") message.delete.side_effect = RuntimeError("missing delete permission") ban_member = AsyncMock() monkeypatch.setattr(service.app, "ban_chat_member", ban_member, raising=False) result = await service.enforce_risk_message( message, { "enabled": True, "keywords": ["做爱"], "actions": ["delete", "ban"], }, ) ban_member.assert_awaited_once() args, kwargs = ban_member.await_args assert args == (-100, 10) assert 3500 <= (kwargs["until_date"] - datetime.now(UTC)).total_seconds() <= 3600 assert result.deleted is False assert result.applied_member_action == "ban" assert result.duration_seconds == 3600 async def test_warning_policy_runs_again_after_threshold_ban_and_unban( app_modules, monkeypatch ): service = app_modules.load("wbb.services.blacklist_enforcement") grant_bot_permissions(app_modules) ban_member = AsyncMock() monkeypatch.setattr(service.app, "ban_chat_member", ban_member, raising=False) policy = { "enabled": True, "keywords": ["做爱"], "actions": ["warn"], "duration_seconds": 1800, } first = await service.enforce_risk_message(make_message(text="做爱"), policy) second = await service.enforce_risk_message(make_message(text="再次做爱"), policy) third = await service.enforce_risk_message(make_message(text="第三次做爱"), policy) after_unban = await service.enforce_risk_message( make_message(text="解封后再次做爱"), policy ) ban_member.assert_awaited_once() args, kwargs = ban_member.await_args assert args == (-100, 10) assert 1700 <= (kwargs["until_date"] - datetime.now(UTC)).total_seconds() <= 1800 assert (first.applied_member_action, first.warning_count) == ("warn", 1) assert (second.applied_member_action, second.warning_count) == ("warn", 2) assert (third.applied_member_action, third.warning_count) == ("ban", 3) assert (after_unban.applied_member_action, after_unban.warning_count) == ( "warn", 1, ) async def test_multiple_matching_rules_merge_delete_and_highest_member_action( app_modules, monkeypatch ): service = app_modules.load("wbb.services.blacklist_enforcement") grant_bot_permissions(app_modules) message = make_message( text="广告请访问 https://example.com", photo=object(), ) ban_member = AsyncMock() unban_member = AsyncMock() monkeypatch.setattr(service.app, "ban_chat_member", ban_member, raising=False) monkeypatch.setattr(service.app, "unban_chat_member", unban_member, raising=False) rules = [ { "rule_id": "keyword-rule", "name": "广告关键词", "enabled": True, "keywords": ["广告"], "actions": ["delete", "warn"], "duration_seconds": 1800, }, { "rule_id": "link-rule", "name": "链接拦截", "enabled": True, "match_links": True, "actions": ["kick"], }, { "rule_id": "image-rule", "name": "图片封禁", "enabled": True, "match_images": True, "actions": ["delete", "ban"], "duration_seconds": 7200, }, ] result = await service.enforce_risk_message(message, rules) message.delete.assert_awaited_once() ban_member.assert_awaited_once() unban_member.assert_not_awaited() assert result.requested_member_action == "ban" assert result.applied_member_action == "ban" assert result.duration_seconds == 7200 assert [match.rule_name for match in result.matches] == [ "广告关键词", "链接拦截", "图片封禁", ] assert all(name in app_modules.app.sent_messages[-1][1] for name in ["广告关键词", "链接拦截", "图片封禁"]) def test_risk_rules_require_unique_ids_valid_triggers_and_ban_duration(app_modules): service = app_modules.load("wbb.services.blacklist_enforcement") base = { "rule_id": "same-id", "name": "测试规则", "enabled": True, "keywords": ["广告"], "actions": ["ban"], "duration_seconds": 3600, } with pytest.raises(service.RiskControlValidationError, match="规则 ID 不能重复"): service.normalize_risk_rules([base, base], strict=True) with pytest.raises(service.RiskControlValidationError, match="触发条件"): service.normalize_risk_rules( [{**base, "rule_id": "empty", "keywords": []}], strict=True ) with pytest.raises(service.RiskControlValidationError, match="处罚时长"): service.normalize_risk_rules( [{**base, "duration_seconds": 30}], strict=True ) async def test_mute_keeps_member_in_chat_and_uses_configured_duration( app_modules, monkeypatch ): service = app_modules.load("wbb.services.blacklist_enforcement") grant_bot_permissions(app_modules) restrict_member = AsyncMock() ban_member = AsyncMock() monkeypatch.setattr(service.app, "restrict_chat_member", restrict_member, raising=False) monkeypatch.setattr(service.app, "ban_chat_member", ban_member, raising=False) result = await service.enforce_risk_message( make_message(text="广告"), [ { "rule_id": "mute-ads", "name": "广告禁言", "enabled": True, "keywords": ["广告"], "actions": ["delete", "mute"], "duration_seconds": 900, } ], ) restrict_member.assert_awaited_once() args, kwargs = restrict_member.await_args assert args[:2] == (-100, 10) assert 800 <= (kwargs["until_date"] - datetime.now(UTC)).total_seconds() <= 900 ban_member.assert_not_awaited() assert result.applied_member_action == "mute" assert result.duration_seconds == 900