fix(message): confine local media attachments

This commit is contained in:
hinotoi-agent 2026-05-15 16:33:43 +08:00 committed by Xubin Ren
parent afbaea870b
commit 57d7847dc8
2 changed files with 95 additions and 14 deletions

View File

@ -1,12 +1,12 @@
"""Message tool for sending messages to users.""" """Message tool for sending messages to users."""
import os
from contextvars import ContextVar from contextvars import ContextVar
from pathlib import Path from pathlib import Path
from typing import Any, Awaitable, Callable from typing import Any, Awaitable, Callable
from nanobot.agent.tools.base import Tool, tool_parameters from nanobot.agent.tools.base import Tool, tool_parameters
from nanobot.agent.tools.context import ContextAware, RequestContext from nanobot.agent.tools.context import ContextAware, RequestContext
from nanobot.agent.tools.filesystem import _resolve_path
from nanobot.agent.tools.schema import ArraySchema, StringSchema, tool_parameters_schema from nanobot.agent.tools.schema import ArraySchema, StringSchema, tool_parameters_schema
from nanobot.bus.events import OutboundMessage from nanobot.bus.events import OutboundMessage
from nanobot.config.paths import get_workspace_path from nanobot.config.paths import get_workspace_path
@ -50,11 +50,19 @@ class MessageTool(Tool, ContextAware):
default_chat_id: str = "", default_chat_id: str = "",
default_message_id: str | None = None, default_message_id: str | None = None,
workspace: str | Path | None = None, workspace: str | Path | None = None,
restrict_to_workspace: bool = False,
): ):
self._send_callback = send_callback self._send_callback = send_callback
self._workspace = Path(workspace).expanduser() if workspace is not None else get_workspace_path() self._workspace = (
self._default_channel: ContextVar[str] = ContextVar("message_default_channel", default=default_channel) Path(workspace).expanduser() if workspace is not None else get_workspace_path()
self._default_chat_id: ContextVar[str] = ContextVar("message_default_chat_id", default=default_chat_id) )
self._restrict_to_workspace = restrict_to_workspace
self._default_channel: ContextVar[str] = ContextVar(
"message_default_channel", default=default_channel
)
self._default_chat_id: ContextVar[str] = ContextVar(
"message_default_chat_id", default=default_chat_id
)
self._default_message_id: ContextVar[str | None] = ContextVar( self._default_message_id: ContextVar[str | None] = ContextVar(
"message_default_message_id", "message_default_message_id",
default=default_message_id, default=default_message_id,
@ -72,7 +80,11 @@ class MessageTool(Tool, ContextAware):
@classmethod @classmethod
def create(cls, ctx: Any) -> Tool: def create(cls, ctx: Any) -> Tool:
send_callback = ctx.bus.publish_outbound if ctx.bus else None send_callback = ctx.bus.publish_outbound if ctx.bus else None
return cls(send_callback=send_callback, workspace=ctx.workspace) return cls(
send_callback=send_callback,
workspace=ctx.workspace,
restrict_to_workspace=ctx.config.restrict_to_workspace,
)
def set_context(self, ctx: RequestContext) -> None: def set_context(self, ctx: RequestContext) -> None:
"""Set the current message context.""" """Set the current message context."""
@ -123,6 +135,20 @@ class MessageTool(Tool, ContextAware):
"Do NOT use read_file to send files — that only reads content for your own analysis." "Do NOT use read_file to send files — that only reads content for your own analysis."
) )
def _resolve_media(self, media: list[str]) -> list[str]:
"""Resolve local media attachments and enforce workspace restriction when enabled."""
resolved: list[str] = []
allowed_dir = self._workspace if self._restrict_to_workspace else None
for p in media:
if p.startswith(("http://", "https://")):
resolved.append(p)
elif not self._restrict_to_workspace:
path = Path(p).expanduser()
resolved.append(p if path.is_absolute() else str(self._workspace / path))
else:
resolved.append(str(_resolve_path(p, self._workspace, allowed_dir)))
return resolved
async def execute( async def execute(
self, self,
content: str, content: str,
@ -131,9 +157,10 @@ class MessageTool(Tool, ContextAware):
message_id: str | None = None, message_id: str | None = None,
media: list[str] | None = None, media: list[str] | None = None,
buttons: list[list[str]] | None = None, buttons: list[list[str]] | None = None,
**kwargs: Any **kwargs: Any,
) -> str: ) -> str:
from nanobot.utils.helpers import strip_think from nanobot.utils.helpers import strip_think
content = strip_think(content) content = strip_think(content)
if buttons is not None: if buttons is not None:
@ -164,13 +191,10 @@ class MessageTool(Tool, ContextAware):
return "Error: Message sending not configured" return "Error: Message sending not configured"
if media: if media:
resolved = [] try:
for p in media: media = self._resolve_media(media)
if p.startswith(("http://", "https://")) or os.path.isabs(p): except (OSError, PermissionError, ValueError) as e:
resolved.append(p) return f"Error: media path is not allowed: {str(e)}"
else:
resolved.append(str(self._workspace / p))
media = resolved
metadata = dict(self._default_metadata.get()) if same_target else {} metadata = dict(self._default_metadata.get()) if same_target else {}
if message_id: if message_id:

View File

@ -30,7 +30,10 @@ async def test_message_tool_rejects_malformed_buttons(bad) -> None:
into the channel layer where Telegram would silently reject the frame.""" into the channel layer where Telegram would silently reject the frame."""
tool = MessageTool() tool = MessageTool()
result = await tool.execute( result = await tool.execute(
content="hi", channel="telegram", chat_id="1", buttons=bad, content="hi",
channel="telegram",
chat_id="1",
buttons=bad,
) )
assert result == "Error: buttons must be a list of list of strings" assert result == "Error: buttons must be a list of list of strings"
@ -84,6 +87,7 @@ async def test_message_tool_inherits_metadata_for_same_target() -> None:
tool = MessageTool(send_callback=_send) tool = MessageTool(send_callback=_send)
slack_meta = {"slack": {"thread_ts": "111.222", "channel_type": "channel"}} slack_meta = {"slack": {"thread_ts": "111.222", "channel_type": "channel"}}
from nanobot.agent.tools.context import RequestContext from nanobot.agent.tools.context import RequestContext
tool.set_context(RequestContext(channel="slack", chat_id="C123", metadata=slack_meta)) tool.set_context(RequestContext(channel="slack", chat_id="C123", metadata=slack_meta))
await tool.execute(content="thread reply") await tool.execute(content="thread reply")
@ -100,6 +104,7 @@ async def test_message_tool_clears_metadata_when_context_has_none() -> None:
tool = MessageTool(send_callback=_send) tool = MessageTool(send_callback=_send)
from nanobot.agent.tools.context import RequestContext from nanobot.agent.tools.context import RequestContext
tool.set_context( tool.set_context(
RequestContext( RequestContext(
channel="slack", channel="slack",
@ -123,6 +128,7 @@ async def test_message_tool_does_not_inherit_metadata_for_cross_target() -> None
tool = MessageTool(send_callback=_send) tool = MessageTool(send_callback=_send)
from nanobot.agent.tools.context import RequestContext from nanobot.agent.tools.context import RequestContext
tool.set_context( tool.set_context(
RequestContext( RequestContext(
channel="slack", channel="slack",
@ -176,6 +182,57 @@ async def test_message_tool_resolves_relative_media_paths_from_active_workspace(
assert sent[0].media == [str(workspace / "output/image.png")] assert sent[0].media == [str(workspace / "output/image.png")]
@pytest.mark.asyncio
async def test_message_tool_rejects_outside_workspace_absolute_media_when_restricted(
tmp_path,
) -> None:
sent: list[OutboundMessage] = []
async def _send(msg: OutboundMessage) -> None:
sent.append(msg)
workspace = tmp_path / "workspace"
workspace.mkdir()
outside = tmp_path / "secret.txt"
outside.write_text("secret", encoding="utf-8")
tool = MessageTool(send_callback=_send, workspace=workspace, restrict_to_workspace=True)
result = await tool.execute(
content="see attached",
channel="telegram",
chat_id="1",
media=[str(outside)],
)
assert result.startswith("Error: media path is not allowed:")
assert "outside allowed directory" in result
assert sent == []
@pytest.mark.asyncio
async def test_message_tool_allows_workspace_absolute_media_when_restricted(tmp_path) -> None:
sent: list[OutboundMessage] = []
async def _send(msg: OutboundMessage) -> None:
sent.append(msg)
workspace = tmp_path / "workspace"
workspace.mkdir()
image = workspace / "image.png"
image.write_text("image", encoding="utf-8")
tool = MessageTool(send_callback=_send, workspace=workspace, restrict_to_workspace=True)
result = await tool.execute(
content="see attached",
channel="telegram",
chat_id="1",
media=[str(image)],
)
assert result == "Message sent to telegram:1 with 1 attachments"
assert sent[0].media == [str(image.resolve())]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_message_tool_passes_through_absolute_media_paths() -> None: async def test_message_tool_passes_through_absolute_media_paths() -> None:
sent: list[OutboundMessage] = [] sent: list[OutboundMessage] = []