shirobana-bot/tests/test_security_channels.py

"""Adversarial tests for command authorization and Telegram sticker delivery."""

import asyncio
import importlib.util
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest import mock

BASE = Path(__file__).resolve().parents[1]
sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]


def load_dispatch():
    """Import dispatch without optional network dependencies of unused modules."""
    stubs = {}
    for name in (
        "setu", "jm", "pixiv", "bilibili", "mail", "live_monitor", "ginkgostu",
        "identity", "identity_cmd", "for_admin_komari_agent", "for_admin_catch_log",
        "mode_commands",
    ):
        stubs[name] = types.ModuleType(name)
    help_ui = types.ModuleType("help_ui")
    help_ui.command_index = lambda: "help"
    stubs["help_ui"] = help_ui
    deepseek = types.ModuleType("deepseek")
    deepseek.__path__ = []
    deepseek.deepseek_balance = types.ModuleType("deepseek.deepseek_balance")
    stubs["deepseek"] = deepseek
    stubs["deepseek.deepseek_balance"] = deepseek.deepseek_balance
    spec = importlib.util.spec_from_file_location(
        "dispatch_security_test", BASE / "bot-utility" / "dispatch.py",
    )
    module = importlib.util.module_from_spec(spec)
    with mock.patch.dict(sys.modules, stubs):
        spec.loader.exec_module(module)
    return module


class ChatAdminTests(unittest.IsolatedAsyncioTestCase):
    async def test_chat_on_off_requires_admin(self):
        dispatch = load_dispatch()
        channel = types.SimpleNamespace(is_admin=mock.Mock(return_value=False))
        engine = types.SimpleNamespace(set_chat_enabled=mock.Mock())
        chat = types.ModuleType("chat")
        chat.engine = engine
        dispatch.send_message = mock.AsyncMock()
        with mock.patch.dict(sys.modules, {"chat": chat}):
            handled = await dispatch._handle_chat_legacy(channel, "group", 20, 10, ".chat on")
        self.assertTrue(handled)
        engine.set_chat_enabled.assert_not_called()
        dispatch.send_message.assert_awaited_once_with(channel, "group", 20, "该指令仅管理员可用")

    async def test_admin_can_toggle_chat(self):
        dispatch = load_dispatch()
        channel = types.SimpleNamespace(is_admin=mock.Mock(return_value=True))
        engine = types.SimpleNamespace(set_chat_enabled=mock.Mock())
        chat = types.ModuleType("chat")
        chat.engine = engine
        dispatch.send_message = mock.AsyncMock()
        with mock.patch.dict(sys.modules, {"chat": chat}):
            for arg, enabled in (("on", True), ("off", False)):
                with self.subTest(arg=arg):
                    handled = await dispatch._handle_chat_legacy(
                        channel, "private", 10, 10, ".chat " + arg,
                    )
                    self.assertTrue(handled)
                    engine.set_chat_enabled.assert_called_with(enabled)
        self.assertEqual(engine.set_chat_enabled.call_count, 2)


try:
    from aiogram.types import FSInputFile
except ModuleNotFoundError:
    aiogram = types.ModuleType("aiogram")
    aiogram.__path__ = []
    aiogram_types = types.ModuleType("aiogram.types")

    class FSInputFile:
        def __init__(self, path):
            self.path = path

    aiogram_types.FSInputFile = FSInputFile
    aiogram.types = aiogram_types
    sys.modules["aiogram"] = aiogram
    sys.modules["aiogram.types"] = aiogram_types

from tg_channel import TGChannel


class TelegramStickerTests(unittest.IsolatedAsyncioTestCase):
    async def test_agent_sticker_uses_sticker_api_not_photo(self):
        with tempfile.TemporaryDirectory() as directory:
            image = Path(directory) / "sticker.webp"
            image.write_bytes(b"sticker")
            bot = types.SimpleNamespace(
                send_sticker=mock.AsyncMock(return_value=object()),
                send_photo=mock.AsyncMock(return_value=object()),
                send_document=mock.AsyncMock(return_value=object()),
            )
            channel = TGChannel(bot, admin_ids=[])
            sent = await channel.send("private", 123, [{
                "type": "image", "data": {"file": image.as_uri(), "sub_type": 1},
            }])
            self.assertTrue(sent)
            bot.send_sticker.assert_awaited_once()
            self.assertEqual(bot.send_sticker.await_args.kwargs["chat_id"], 123)
            bot.send_photo.assert_not_awaited()
            bot.send_document.assert_not_awaited()

    async def test_regular_image_still_uses_photo_api(self):
        with tempfile.TemporaryDirectory() as directory:
            image = Path(directory) / "image.jpg"
            image.write_bytes(b"image")
            bot = types.SimpleNamespace(
                send_sticker=mock.AsyncMock(return_value=object()),
                send_photo=mock.AsyncMock(return_value=object()),
                send_document=mock.AsyncMock(return_value=object()),
            )
            channel = TGChannel(bot, admin_ids=[])
            sent = await channel.send("private", 123, [{
                "type": "image", "data": {"file": image.as_uri()},
            }])
            self.assertTrue(sent)
            bot.send_photo.assert_awaited_once()
            bot.send_sticker.assert_not_awaited()


if __name__ == "__main__":
    unittest.main()