diff --git a/nanobot/agent/tools/message.py b/nanobot/agent/tools/message.py index 339f9bdcf..54a196e40 100644 --- a/nanobot/agent/tools/message.py +++ b/nanobot/agent/tools/message.py @@ -1,12 +1,12 @@ """Message tool for sending messages to users.""" -import os from contextvars import ContextVar from pathlib import Path from typing import Any, Awaitable, Callable from nanobot.agent.tools.base import Tool, tool_parameters 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.bus.events import OutboundMessage from nanobot.config.paths import get_workspace_path @@ -50,11 +50,19 @@ class MessageTool(Tool, ContextAware): default_chat_id: str = "", default_message_id: str | None = None, workspace: str | Path | None = None, + restrict_to_workspace: bool = False, ): self._send_callback = send_callback - self._workspace = Path(workspace).expanduser() if workspace is not None else get_workspace_path() - 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._workspace = ( + Path(workspace).expanduser() if workspace is not None else get_workspace_path() + ) + 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( "message_default_message_id", default=default_message_id, @@ -72,7 +80,11 @@ class MessageTool(Tool, ContextAware): @classmethod def create(cls, ctx: Any) -> Tool: 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: """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." ) + 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( self, content: str, @@ -131,9 +157,10 @@ class MessageTool(Tool, ContextAware): message_id: str | None = None, media: list[str] | None = None, buttons: list[list[str]] | None = None, - **kwargs: Any + **kwargs: Any, ) -> str: from nanobot.utils.helpers import strip_think + content = strip_think(content) if buttons is not None: @@ -164,13 +191,10 @@ class MessageTool(Tool, ContextAware): return "Error: Message sending not configured" if media: - resolved = [] - for p in media: - if p.startswith(("http://", "https://")) or os.path.isabs(p): - resolved.append(p) - else: - resolved.append(str(self._workspace / p)) - media = resolved + try: + media = self._resolve_media(media) + except (OSError, PermissionError, ValueError) as e: + return f"Error: media path is not allowed: {str(e)}" metadata = dict(self._default_metadata.get()) if same_target else {} if message_id: diff --git a/tests/tools/test_message_tool.py b/tests/tools/test_message_tool.py index d4439422a..fc37217a2 100644 --- a/tests/tools/test_message_tool.py +++ b/tests/tools/test_message_tool.py @@ -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.""" tool = MessageTool() 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" @@ -84,6 +87,7 @@ async def test_message_tool_inherits_metadata_for_same_target() -> None: tool = MessageTool(send_callback=_send) slack_meta = {"slack": {"thread_ts": "111.222", "channel_type": "channel"}} from nanobot.agent.tools.context import RequestContext + tool.set_context(RequestContext(channel="slack", chat_id="C123", metadata=slack_meta)) 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) from nanobot.agent.tools.context import RequestContext + tool.set_context( RequestContext( channel="slack", @@ -123,6 +128,7 @@ async def test_message_tool_does_not_inherit_metadata_for_cross_target() -> None tool = MessageTool(send_callback=_send) from nanobot.agent.tools.context import RequestContext + tool.set_context( RequestContext( 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")] +@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 async def test_message_tool_passes_through_absolute_media_paths() -> None: sent: list[OutboundMessage] = []