TelegramBotAdmin/tests/test_blacklist_enforcement.py

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("ＳＰＡＭ 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