Initial import of Shirobana bot backend

Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
2026-11-11 00:25:30 +08:00
commit 34089d0042
106 changed files with 28028 additions and 0 deletions
@@ -0,0 +1,417 @@
+"""Isolated tests for the agent tool runtime; no real API or message delivery."""
+
+import asyncio
+import base64
+import json
+import os
+import shlex
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest import mock
+
+BASE = Path(__file__).resolve().parents[1]
+sys.path[:0] = [str(BASE / "bot-utility"), str(BASE)]
+
+from agent import ArtifactError, ArtifactStore, CaptureChannel, ModeStore, ToolContext, ToolExecutor
+import channel as channel_module
+
+
+class FakeChannel:
+    name = "qq"
+
+    def __init__(self):
+        self.sent = []
+
+    async def send(self, msg_type, target_id, content):
+        self.sent.append((msg_type, target_id, content))
+        return True
+
+
+class ModeStoreTests(unittest.TestCase):
+    def test_default_private_and_group_scopes(self):
+        with tempfile.TemporaryDirectory() as directory:
+            store = ModeStore(Path(directory) / "agent.db")
+            self.assertEqual(store.get_mode("qq", 10), "command")
+            store.set_mode("qq", 10, "agent")
+            store.set_mode("qq", 10, "agent", group_id=20)
+            store.set_mode("tg", 10, "agent")
+            self.assertEqual(store.get_mode("qq", 10), "agent")
+            self.assertEqual(store.get_mode("qq", 999, group_id=20), "agent")
+            self.assertEqual(store.get_mode("qq", 10, group_id=21), "command")
+            store.set_mode("qq", 10, "command", group_id=20)
+            self.assertEqual(store.get_mode("qq", 10, group_id=20), "command")
+            with self.assertRaises(ValueError):
+                store.set_mode("qq", 10, "invalid")
+
+
+class ArtifactStoreTests(unittest.TestCase):
+    def test_local_base64_and_confined_materialize(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory) / "artifacts"
+            source = Path(directory) / "source.jpg"
+            source.write_bytes(b"image-data")
+            store = ArtifactStore(root)
+            local = store.register(source, "image")
+            inline = store.register(base64.b64encode(b"second").decode(), "image", "two.jpg")
+            self.assertNotEqual(local["artifact_id"], inline["artifact_id"])
+            self.assertGreaterEqual(len(local["artifact_id"]), 16)
+            first_path = store.materialize(local["artifact_id"])
+            self.assertEqual(first_path.read_bytes(), b"image-data")
+            self.assertTrue(first_path.is_relative_to(root))
+            source.unlink()
+            self.assertEqual(first_path.read_bytes(), b"image-data")
+            self.assertNotIn(str(source), json.dumps(local))
+            with self.assertRaises(ArtifactError):
+                store.materialize("../../etc/passwd")
+            with self.assertRaises(ArtifactError):
+                store.register("/definitely/not/a/real/file")
+
+    def test_url_materializes_under_root(self):
+        with tempfile.TemporaryDirectory() as directory:
+            store = ArtifactStore(Path(directory) / "store")
+            artifact = store.register("https://example.invalid/a.jpg", "image")
+            response = mock.MagicMock()
+            response.__enter__.return_value = response
+            response.geturl.return_value = "https://example.invalid/a.jpg"
+            response.headers = {}
+            response.read.return_value = b"remote-image"
+            with mock.patch("urllib.request.urlopen", return_value=response) as opener:
+                path = store.materialize(artifact["artifact_id"])
+            self.assertEqual(path.read_bytes(), b"remote-image")
+            self.assertTrue(path.is_relative_to(store.root))
+            opener.assert_called_once()
+
+    def test_scoped_artifacts_require_matching_owner_and_session(self):
+        with tempfile.TemporaryDirectory() as directory:
+            store = ArtifactStore(Path(directory) / "store")
+            channel = SimpleNamespace(name="qq")
+            ctx = SimpleNamespace(channel=channel, msg_type="group", user_id=10, target_id=20)
+            artifact = store.register(b64("private"), "file", "private.txt", ctx=ctx)
+            artifact_id = artifact["artifact_id"]
+            self.assertEqual(store.get(artifact_id, ctx)["artifact_id"], artifact_id)
+            self.assertEqual(store.materialize(artifact_id, ctx).read_bytes(), b"private")
+            record = json.loads((store.root / artifact_id / "artifact.json").read_text())
+            self.assertEqual(record["owner_id"], "10")
+            self.assertEqual(record["session"], {
+                "channel": "qq", "msg_type": "group", "target_id": "20",
+            })
+            foreign_contexts = [
+                SimpleNamespace(channel=channel, msg_type="group", user_id=11, target_id=20),
+                SimpleNamespace(channel=channel, msg_type="group", user_id=10, target_id=21),
+                SimpleNamespace(channel=channel, msg_type="private", user_id=10, target_id=20),
+                SimpleNamespace(channel=SimpleNamespace(name="tg"), msg_type="group", user_id=10, target_id=20),
+            ]
+            for foreign in foreign_contexts:
+                with self.subTest(foreign=foreign):
+                    with self.assertRaises(ArtifactError):
+                        store.get(artifact_id, foreign)
+                    with self.assertRaises(ArtifactError):
+                        store.materialize(artifact_id, foreign)
+            with self.assertRaises(ArtifactError):
+                store.get(artifact_id)
+            with self.assertRaises(ArtifactError):
+                store.materialize(artifact_id)
+
+    def test_legacy_artifacts_are_denied_to_agent_contexts(self):
+        with tempfile.TemporaryDirectory() as directory:
+            store = ArtifactStore(Path(directory) / "store")
+            legacy = store.register(b64("legacy"), "file")
+            ctx = SimpleNamespace(
+                channel=SimpleNamespace(name="qq"), msg_type="private",
+                user_id=10, target_id=10,
+            )
+            artifact_id = legacy["artifact_id"]
+            self.assertEqual(store.materialize(artifact_id).read_bytes(), b"legacy")
+            with self.assertRaises(ArtifactError):
+                store.get(artifact_id, ctx)
+            with self.assertRaises(ArtifactError):
+                store.materialize(artifact_id, ctx)
+
+
+class ExecutorTests(unittest.IsolatedAsyncioTestCase):
+    def setUp(self):
+        self.temporary = tempfile.TemporaryDirectory()
+        self.root = Path(self.temporary.name)
+        self.store = ArtifactStore(self.root / "artifacts")
+        self.real_channel = FakeChannel()
+        self.sender = mock.AsyncMock(return_value=True)
+        self.ctx = ToolContext(
+            self.real_channel, "private", 9000000001, 123,
+            is_admin=False, request_id="req", task_id="task",
+        )
+
+    def tearDown(self):
+        self.temporary.cleanup()
+
+    def executor(self, modules):
+        return ToolExecutor(self.store, modules=modules, send_message=self.sender)
+
+    async def test_capture_channel_never_delivers_or_mirrors_private_text(self):
+        capture = CaptureChannel(self.ctx, self.store)
+        # Use the production sender to prove its private mirror path is skipped.
+        with mock.patch.object(channel_module, "_mirror_private", new=mock.AsyncMock()) as mirror:
+            result = await channel_module.send_message(capture, "private", 123, "captured")
+        self.assertFalse(result)
+        self.assertEqual(capture.texts, ["captured"])
+        self.assertEqual(self.real_channel.sent, [])
+        mirror.assert_not_awaited()
+
+    async def test_setu_media_is_delivered_once_and_text_is_model_facing(self):
+        image = self.root / "captured.jpg"
+        image.write_bytes(b"jpeg")
+        calls = []
+
+        async def handle_setu(tags, r18, ai, msg_type, target_id, ws):
+            from channel import send_message
+            calls.append((tags, r18, ai))
+            await send_message(ws, msg_type, target_id, "正在获取图片")
+            await send_message(ws, msg_type, target_id, [
+                {"type": "text", "data": {"text": "PID 42"}},
+                {"type": "image", "data": {"file": str(image)}},
+            ])
+
+        module = SimpleNamespace(handle_setu=handle_setu)
+        executor = self.executor({"setu": module})
+        result = await executor.execute(
+            "setu", {"action": "get", "params": {"tags": ["猫耳"], "r18": False}}, self.ctx,
+        )
+        self.assertTrue(result.ok, result)
+        self.assertEqual(calls, [(["猫耳"], "不要r18", None)])
+        self.assertEqual(result.safe_message, "正在获取图片\nPID 42")
+        self.assertEqual(len(result.artifacts), 1)
+        self.assertEqual(result.data["deliveries"][0]["ok"], True)
+        self.sender.assert_awaited_once()
+        sent = self.sender.await_args.args[3]
+        self.assertEqual(sent[0]["type"], "image")
+        delivered = Path(sent[0]["data"]["file"].removeprefix("file://"))
+        self.assertTrue(delivered.is_relative_to(self.store.root))
+        self.assertEqual(self.real_channel.sent, [])
+
+    async def test_other_tool_text_is_not_delivered_and_mail_quotes_are_preserved(self):
+        seen = []
+
+        async def handle_mail(ws, msg_type, target_id, user_id, raw):
+            from channel import send_message
+            seen.append(raw)
+            await send_message(ws, msg_type, target_id, "邮件已发送")
+
+        executor = self.executor({"mail": SimpleNamespace(handle_mail=handle_mail)})
+        result = await executor.execute("mail", {
+            "action": "new",
+            "params": {"to_address": "to@example.com", "subject": "two words", "body_text": "hello world"},
+        }, self.ctx)
+        self.assertTrue(result.ok)
+        tokens = shlex.split(seen[0])
+        self.assertEqual(tokens[-3:], ["to@example.com", "two words", "hello world"])
+        self.assertEqual(result.safe_message, "邮件已发送")
+        self.sender.assert_not_awaited()
+        self.assertEqual(self.real_channel.sent, [])
+
+    async def test_mail_send_unknown_is_not_retryable(self):
+        async def handle_mail(ws, msg_type, target_id, user_id, raw):
+            from channel import send_message
+            await send_message(ws, msg_type, target_id, "发送失败：上游返回了无法识别的响应")
+
+        executor = self.executor({"mail": SimpleNamespace(handle_mail=handle_mail)})
+        result = await executor.execute("mail", {
+            "action": "send",
+            "params": {"to_address": "to@example.com", "subject": "s", "body_text": "b"},
+        }, self.ctx)
+        self.assertFalse(result.ok)
+        self.assertEqual(result.code, "send_unknown")
+        self.assertFalse(result.retryable)
+
+    async def test_outbound_rejects_artifact_from_another_owner(self):
+        foreign_ctx = ToolContext(self.real_channel, "private", self.ctx.user_id + 1, 123)
+        foreign = self.store.register(b64("secret"), "file", "secret.txt", ctx=foreign_ctx)
+        executor = self.executor({})
+        result = await executor.execute(
+            "send_file", {"artifact_id": foreign["artifact_id"]}, self.ctx,
+        )
+        self.assertFalse(result.ok)
+        self.assertEqual(result.code, "invalid_artifact")
+        self.sender.assert_not_awaited()
+        self.assertEqual(self.real_channel.sent, [])
+
+    async def test_mail_attachments_reject_foreign_artifact_before_sending(self):
+        foreign_ctx = ToolContext(self.real_channel, "private", self.ctx.user_id, 999)
+        foreign = self.store.register(b64("secret"), "file", "secret.txt", ctx=foreign_ctx)
+        send_with_attachments = mock.AsyncMock(return_value={"ok": True})
+        module = SimpleNamespace(
+            _load_user=lambda uid: {
+                "default": "from@rrrosmontis.icu",
+                "addresses": ["from@rrrosmontis.icu"],
+            },
+            send_with_attachments=send_with_attachments,
+        )
+        executor = self.executor({"mail": module})
+        result = await executor.execute("mail", {
+            "action": "send",
+            "params": {
+                "to_address": "to@example.com", "subject": "s", "body_text": "b",
+                "artifact_ids": [foreign["artifact_id"]],
+            },
+        }, self.ctx)
+        self.assertFalse(result.ok)
+        self.assertEqual(result.code, "invalid_artifact")
+        send_with_attachments.assert_not_awaited()
+
+    async def test_mail_list_limit_and_info_mid_aliases(self):
+        seen = []
+        async def handle_mail(ws, msg_type, target_id, user_id, raw):
+            seen.append(raw)
+        executor = self.executor({"mail": SimpleNamespace(handle_mail=handle_mail)})
+        result = await executor.execute("mail", {"action": "list", "params": {"limit": 5}}, self.ctx)
+        self.assertTrue(result.ok, result)
+        result = await executor.execute("mail", {"action": "info", "params": {"mid": 2645}}, self.ctx)
+        self.assertTrue(result.ok, result)
+        self.assertEqual(seen, [".mail list 5", ".mail info 2645"])
+
+    async def test_mail_attachments_use_new_api_and_only_artifact_paths(self):
+        artifact = self.store.register(b64("attachment"), "file", "a.txt", ctx=self.ctx)
+        calls = []
+
+        async def send_with_attachments(*args):
+            calls.append(args)
+            return {"ok": True, "attachment_count": 1}
+
+        async def handle_mail(*args):
+            raise AssertionError("legacy mail handler must not be called")
+
+        module = SimpleNamespace(
+            _load_user=lambda uid: {
+                "default": "from@rrrosmontis.icu",
+                "addresses": ["from@rrrosmontis.icu"],
+            },
+            send_with_attachments=send_with_attachments,
+            handle_mail=handle_mail,
+        )
+        executor = self.executor({"mail": module})
+        result = await executor.execute("mail", {
+            "action": "send",
+            "params": {
+                "to_address": "to@example.com", "subject": "s", "body_text": "b",
+                "artifact_ids": [artifact["artifact_id"]],
+            },
+        }, self.ctx)
+        self.assertTrue(result.ok, result)
+        self.assertEqual(result.code, "sent")
+        self.assertEqual(len(calls), 1)
+        args = calls[0]
+        self.assertEqual(args[:5], (
+            self.ctx.user_id, "from@rrrosmontis.icu", "to@example.com", "s", "b",
+        ))
+        self.assertEqual(Path(args[5][0]).read_bytes(), b"attachment")
+        self.assertTrue(Path(args[5][0]).is_relative_to(self.store.root))
+        self.sender.assert_not_awaited()
+
+    async def test_permission_is_enforced_before_admin_handler_import_or_call(self):
+        handler = mock.AsyncMock()
+        executor = self.executor({"server": SimpleNamespace(handle_command=handler)})
+        result = await executor.execute("server", {"action": "list", "params": {}}, self.ctx)
+        self.assertFalse(result.ok)
+        self.assertEqual(result.code, "permission_denied")
+        handler.assert_not_awaited()
+        schema_names = {item["name"] for item in executor.prompt_schema(self.ctx)}
+        self.assertNotIn("server", schema_names)
+        admin_ctx = ToolContext(self.real_channel, "group", 1, 2, is_admin=True)
+        self.assertIn("server", {item["name"] for item in executor.prompt_schema(admin_ctx)})
+
+    async def test_send_file_accepts_only_artifact_id(self):
+        executor = self.executor({})
+        rejected = await executor.execute("send_file", {"artifact_id": "/etc/passwd"}, self.ctx)
+        self.assertFalse(rejected.ok)
+        self.assertEqual(rejected.code, "invalid_artifact")
+        self.sender.assert_not_awaited()
+        artifact = self.store.register(b64("file-body"), "file", "report.txt", ctx=self.ctx)
+        result = await executor.execute("send_file", {"artifact_id": artifact["artifact_id"]}, self.ctx)
+        self.assertTrue(result.ok, result)
+        content = self.sender.await_args.args[3]
+        self.assertEqual(content[0]["type"], "file")
+        self.assertNotEqual(content[0]["data"]["file"], "/etc/passwd")
+        self.assertTrue(Path(content[0]["data"]["file"].removeprefix("file://")).is_relative_to(self.store.root))
+
+    async def test_send_sticker_keeps_onebot_image_segment_for_qq(self):
+        artifact = self.store.register(b64("sticker"), "image", "sticker.webp", ctx=self.ctx)
+        executor = self.executor({})
+        result = await executor.execute(
+            "send_sticker", {"artifact_id": artifact["artifact_id"]}, self.ctx,
+        )
+        self.assertTrue(result.ok, result)
+        content = self.sender.await_args.args[3]
+        self.assertEqual(content[0]["type"], "image")
+        self.assertEqual(content[0]["data"]["sub_type"], 1)
+
+    async def test_validation_and_handler_errors_are_safe(self):
+        async def broken(*args):
+            raise RuntimeError("secret=/home/private/token")
+
+        executor = self.executor({"log": SimpleNamespace(handle_command=broken)})
+        admin = ToolContext(self.real_channel, "private", 1, 1, is_admin=True)
+        result = await executor.execute("log", {"action": "tail", "params": {"lines": 5}}, admin)
+        self.assertFalse(result.ok)
+        self.assertTrue(result.retryable)
+        self.assertNotIn("secret", result.safe_message)
+        bad = await executor.execute("log", {"action": "tail", "params": {"lines": "5;rm"}}, admin)
+        self.assertEqual(bad.code, "invalid_args")
+        unknown = await executor.execute("web_search", {}, admin)
+        self.assertEqual(unknown.code, "unknown_tool")
+
+    async def test_remaining_legacy_modules_are_awaited_with_safe_commands(self):
+        calls = []
+
+        def module(handler_name):
+            async def handler(*args):
+                calls.append((handler_name, args[-1]))
+            return SimpleNamespace(**{handler_name: handler})
+
+        live = SimpleNamespace(
+            HELP_TEXT="live help",
+            status_text=lambda *args: "status",
+            add_follow=lambda *args: calls.append(("add_follow", args)),
+            remove_follow=lambda *args: None,
+        )
+        modules = {
+            "pixiv": module("handle_pixiv"),
+            "jm": module("handle_jm"),
+            "ginkgo": module("handle_ginkgo"),
+            "my": module("handle_my"),
+            "deepseek": module("handle_command"),
+            "log": module("handle_command"),
+            "live": live,
+        }
+        executor = self.executor(modules)
+        admin = ToolContext(self.real_channel, "group", 1, 2, is_admin=True)
+        requests = [
+            ("pixiv", "get", {"pid": 123}),
+            ("jm", "search", {"keyword": "two words", "page": 2}),
+            ("ginkgo", "notice", {"count": 3}),
+            ("my", "push", {"platform": "qq", "enabled": True}),
+            ("deepseek", "alert", {"enabled": False}),
+            ("log", "tail", {"lines": 7}),
+            ("live", "follow", {"room": 4, "site": 1}),
+        ]
+        for name, action, params in requests:
+            result = await executor.execute(name, {"action": action, "params": params}, admin)
+            self.assertTrue(result.ok, (name, result))
+        self.assertEqual(calls[0], ("handle_pixiv", ".pixiv 123"))
+        self.assertEqual(calls[1], ("handle_jm", ".jm ss two words 2"))
+        self.assertEqual(calls[2], ("handle_ginkgo", ".ginkgo notice 3"))
+        self.assertEqual(calls[3], ("handle_my", "qq-push on"))
+        self.assertEqual(calls[4], ("handle_command", ".deepseek alert off"))
+        self.assertEqual(calls[5], ("handle_command", ".log 7"))
+        self.assertEqual(calls[6][0], "add_follow")
+        self.assertEqual(self.real_channel.sent, [])
+        self.sender.assert_not_awaited()
+
+
+def b64(value):
+    """Encode test payload without embedding large base64 literals."""
+    return base64.b64encode(value.encode()).decode()
+
+
+if __name__ == "__main__":
+    unittest.main()