mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-05-20 00:22:31 +00:00
fix(message): confine local media attachments
This commit is contained in:
parent
afbaea870b
commit
57d7847dc8
@ -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:
|
||||||
|
|||||||
@ -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] = []
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user