mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
* fix: restrict local file paths in message tool * Update astrbot/core/tools/message_tools.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * fix: rf --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
238 lines
8.3 KiB
Python
238 lines
8.3 KiB
Python
"""Tests for send_message_to_user session handling."""
|
||
|
||
from types import SimpleNamespace
|
||
from unittest.mock import AsyncMock
|
||
|
||
import pytest
|
||
|
||
from astrbot.core.tools.message_tools import SendMessageToUserTool
|
||
|
||
|
||
def _make_context(
|
||
current_session="feishu:GroupMessage:oc_xxx",
|
||
role="admin",
|
||
require_admin=True,
|
||
runtime="local",
|
||
):
|
||
"""Build a minimal ContextWrapper for SendMessageToUserTool."""
|
||
cfg = {
|
||
"provider_settings": {
|
||
"computer_use_require_admin": require_admin,
|
||
"computer_use_runtime": runtime,
|
||
}
|
||
}
|
||
return SimpleNamespace(
|
||
context=SimpleNamespace(
|
||
event=SimpleNamespace(
|
||
unified_msg_origin=current_session,
|
||
role=role,
|
||
get_sender_id=lambda: "user-1",
|
||
),
|
||
context=SimpleNamespace(
|
||
get_config=lambda umo: cfg,
|
||
send_message=AsyncMock(),
|
||
),
|
||
)
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_send_message_with_full_three_part_session():
|
||
"""LLM passes a complete three-part session string."""
|
||
tool = SendMessageToUserTool()
|
||
ctx = _make_context(current_session="feishu:GroupMessage:oc_aaa")
|
||
result = await tool.call(
|
||
ctx,
|
||
messages=[{"type": "plain", "text": "hello"}],
|
||
session="feishu:GroupMessage:oc_aaa",
|
||
)
|
||
assert "Message sent to session" in result
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_send_message_with_partial_session_id_fallback():
|
||
"""LLM passes only session_id (no colons) — fallback to current_session's prefix."""
|
||
tool = SendMessageToUserTool()
|
||
ctx = _make_context(current_session="feishu:GroupMessage:oc_abc")
|
||
result = await tool.call(
|
||
ctx,
|
||
messages=[{"type": "plain", "text": "hello"}],
|
||
session="oc_abc",
|
||
)
|
||
assert "Message sent to session" in result
|
||
# Verify the target session was reconstructed with current_session's platform/msg_type
|
||
call_args = ctx.context.context.send_message.call_args
|
||
target_session = call_args[0][0]
|
||
assert target_session.platform_id == "feishu"
|
||
assert target_session.message_type.value == "GroupMessage"
|
||
assert target_session.session_id == "oc_abc"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_send_message_defaults_to_current_session():
|
||
"""LLM does not pass session — uses current_session directly."""
|
||
tool = SendMessageToUserTool()
|
||
ctx = _make_context(current_session="feishu:GroupMessage:oc_xxx")
|
||
result = await tool.call(
|
||
ctx,
|
||
messages=[{"type": "plain", "text": "hello"}],
|
||
)
|
||
assert "Message sent to session" in result
|
||
call_args = ctx.context.context.send_message.call_args
|
||
target_session = call_args[0][0]
|
||
assert str(target_session) == "feishu:GroupMessage:oc_xxx"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_send_message_partial_session_falls_back_to_current():
|
||
"""LLM passes session_id matching current_session's id — same session, just incomplete format."""
|
||
tool = SendMessageToUserTool()
|
||
ctx = _make_context(current_session="qq_official:GroupMessage:g123")
|
||
result = await tool.call(
|
||
ctx,
|
||
messages=[{"type": "plain", "text": "world"}],
|
||
session="g123",
|
||
)
|
||
assert "Message sent to session" in result
|
||
call_args = ctx.context.context.send_message.call_args
|
||
target_session = call_args[0][0]
|
||
assert target_session.platform_id == "qq_official"
|
||
assert target_session.message_type.value == "GroupMessage"
|
||
assert target_session.session_id == "g123"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_cron_context_current_session_is_target_session():
|
||
"""在 cron 场景中,current_session 就是 cron 任务的目标 session。
|
||
|
||
cron 是主动唤醒,没有用户消息触发,因此没有"正在聊天的 session"。
|
||
event.unified_msg_origin 来自 CronMessageEvent.session,
|
||
而 CronMessageEvent.session 来自 cron job payload.session,
|
||
即用户在 cron 配置中填写的目标会话。
|
||
"""
|
||
tool = SendMessageToUserTool()
|
||
# cron 任务的目标 session(用户配置的完整三段式)
|
||
cron_target_session = "feishu:GroupMessage:oc_cron_target"
|
||
ctx = _make_context(current_session=cron_target_session)
|
||
|
||
# LLM 在 cron 上下文中只传了 session_id 部分
|
||
result = await tool.call(
|
||
ctx,
|
||
messages=[{"type": "plain", "text": "cron message"}],
|
||
session="oc_cron_target",
|
||
)
|
||
assert "Message sent to session" in result
|
||
call_args = ctx.context.context.send_message.call_args
|
||
target_session = call_args[0][0]
|
||
# 补全后的 session 应与 cron 目标 session 完全一致
|
||
assert str(target_session) == cron_target_session
|
||
assert target_session.platform_id == "feishu"
|
||
assert target_session.message_type.value == "GroupMessage"
|
||
assert target_session.session_id == "oc_cron_target"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_send_message_empty_messages_returns_error():
|
||
"""Empty or missing messages returns error before session resolution."""
|
||
tool = SendMessageToUserTool()
|
||
ctx = _make_context()
|
||
result = await tool.call(ctx, messages=[], session="oc_xxx")
|
||
assert "error:" in result
|
||
assert "messages" in result.lower()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_send_message_missing_image_path_stops_before_send(tmp_path, monkeypatch):
|
||
"""Missing image paths fail before sending any message components."""
|
||
tool = SendMessageToUserTool()
|
||
ctx = _make_context()
|
||
missing_image_path = tmp_path / "missing.png"
|
||
|
||
async def mock_get_booter(*args, **kwargs):
|
||
del args, kwargs
|
||
raise RuntimeError("sandbox unavailable")
|
||
|
||
monkeypatch.setattr(
|
||
"astrbot.core.tools.message_tools.get_booter",
|
||
mock_get_booter,
|
||
)
|
||
|
||
result = await tool.call(
|
||
ctx,
|
||
messages=[
|
||
{"type": "plain", "text": "before image"},
|
||
{"type": "image", "path": str(missing_image_path)},
|
||
],
|
||
)
|
||
|
||
assert "error: failed to build messages[1] component: sandbox unavailable" in result
|
||
ctx.context.context.send_message.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_non_admin_cannot_send_arbitrary_local_absolute_file(tmp_path):
|
||
"""Non-admin users cannot send host files outside the allowed local roots."""
|
||
tool = SendMessageToUserTool()
|
||
ctx = _make_context(role="member", require_admin=True)
|
||
secret_path = tmp_path / "secret.txt"
|
||
secret_path.write_text("secret", encoding="utf-8")
|
||
|
||
result = await tool.call(
|
||
ctx,
|
||
messages=[{"type": "file", "path": str(secret_path)}],
|
||
)
|
||
|
||
assert "error: Local file send is restricted for this user" in result
|
||
assert str(secret_path) in result
|
||
ctx.context.context.send_message.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_non_admin_can_send_workspace_file(tmp_path, monkeypatch):
|
||
"""Non-admin users can send files inside their per-session workspace."""
|
||
tool = SendMessageToUserTool()
|
||
ctx = _make_context(
|
||
current_session="feishu:GroupMessage:oc_workspace",
|
||
role="member",
|
||
require_admin=True,
|
||
)
|
||
workspace_root = tmp_path / "workspaces"
|
||
workspace_file = workspace_root / "feishu_GroupMessage_oc_workspace" / "result.txt"
|
||
workspace_file.parent.mkdir(parents=True)
|
||
workspace_file.write_text("result", encoding="utf-8")
|
||
monkeypatch.setattr(
|
||
"astrbot.core.tools.computer_tools.util.get_astrbot_workspaces_path",
|
||
lambda: str(workspace_root),
|
||
)
|
||
|
||
result = await tool.call(
|
||
ctx,
|
||
messages=[{"type": "file", "path": "result.txt"}],
|
||
)
|
||
|
||
assert "Message sent to session" in result
|
||
ctx.context.context.send_message.assert_called_once()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_non_admin_can_send_temp_file(tmp_path, monkeypatch):
|
||
"""Non-admin users can send generated files under AstrBot temp."""
|
||
tool = SendMessageToUserTool()
|
||
ctx = _make_context(role="member", require_admin=True)
|
||
temp_root = tmp_path / "temp"
|
||
temp_root.mkdir()
|
||
output_path = temp_root / "output.txt"
|
||
output_path.write_text("output", encoding="utf-8")
|
||
monkeypatch.setattr(
|
||
"astrbot.core.tools.message_tools.get_astrbot_temp_path",
|
||
lambda: str(temp_root),
|
||
)
|
||
|
||
result = await tool.call(
|
||
ctx,
|
||
messages=[{"type": "file", "path": str(output_path)}],
|
||
)
|
||
|
||
assert "Message sent to session" in result
|
||
ctx.context.context.send_message.assert_called_once()
|