mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
2472a12671
* feat: filesystem grep, read, edit file * feat: add file write tool and enhance file read functionality * feat: enhance tool prompt formatting and add escaped text decoding for file editing * feat: remove redundant safe path tests from security restrictions * feat: implement file read tool with support for text and image files, including validation for large files * feat: add file read utilities and integrate with filesystem tools * refactor: move computer tools to builtin tools registry * refactor: remove unused plugin_context parameter from _apply_sandbox_tools * feat: supports to display enabled builtin tools in configs * feat: add tooltip for disabled builtin tools and update localization strings * feat: add workspace extra prompt handling in message processing * feat: add ripgrep installation to Dockerfile * perf: shell executed in workspace dir in local env * feat: enhance file reading capabilities to support PDF and DOCX parsing, including workspace storage for long documents * feat: update converted text notice to suggest using grep for large files * feat: implement handling for large tool results with overflow file writing and read tool integration * fix: test * feat: enhance onboarding steps to include computer access configuration and related help information * feat: add support for additional temporary path in restricted environment checks * feat: update computer access hints and add detailed configuration instructions
343 lines
11 KiB
Python
343 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import locale
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from python_ripgrep import search
|
|
|
|
from astrbot.api import logger
|
|
from astrbot.core.computer.file_read_utils import (
|
|
detect_text_encoding,
|
|
read_local_text_range_sync,
|
|
)
|
|
from astrbot.core.utils.astrbot_path import get_astrbot_root
|
|
|
|
from ..olayer import FileSystemComponent, PythonComponent, ShellComponent
|
|
from .base import ComputerBooter
|
|
from .shipyard_search_file_util import _truncate_long_lines
|
|
|
|
_BLOCKED_COMMAND_PATTERNS = [
|
|
" rm -rf ",
|
|
" rm -fr ",
|
|
" rm -r ",
|
|
" mkfs",
|
|
" dd if=",
|
|
" shutdown",
|
|
" reboot",
|
|
" poweroff",
|
|
" halt",
|
|
" sudo ",
|
|
":(){:|:&};:",
|
|
" kill -9 ",
|
|
" killall ",
|
|
]
|
|
|
|
|
|
def _is_safe_command(command: str) -> bool:
|
|
cmd = f" {command.strip().lower()} "
|
|
return not any(pat in cmd for pat in _BLOCKED_COMMAND_PATTERNS)
|
|
|
|
|
|
def _decode_bytes_with_fallback(
|
|
output: bytes | None,
|
|
*,
|
|
preferred_encoding: str | None = None,
|
|
) -> str:
|
|
if output is None:
|
|
return ""
|
|
|
|
preferred = locale.getpreferredencoding(False) or "utf-8"
|
|
attempted_encodings: list[str] = []
|
|
|
|
def _try_decode(encoding: str) -> str | None:
|
|
normalized = encoding.lower()
|
|
if normalized in attempted_encodings:
|
|
return None
|
|
attempted_encodings.append(normalized)
|
|
try:
|
|
return output.decode(encoding)
|
|
except (LookupError, UnicodeDecodeError):
|
|
return None
|
|
|
|
for encoding in filter(None, [preferred_encoding, "utf-8", "utf-8-sig"]):
|
|
if decoded := _try_decode(encoding):
|
|
return decoded
|
|
|
|
if os.name == "nt":
|
|
for encoding in ("mbcs", "cp936", "gbk", "gb18030", preferred):
|
|
if decoded := _try_decode(encoding):
|
|
return decoded
|
|
elif decoded := _try_decode(preferred):
|
|
return decoded
|
|
|
|
return output.decode("utf-8", errors="replace")
|
|
|
|
|
|
def _decode_shell_output(output: bytes | None) -> str:
|
|
return _decode_bytes_with_fallback(output, preferred_encoding="utf-8")
|
|
|
|
|
|
@dataclass
|
|
class LocalShellComponent(ShellComponent):
|
|
async def exec(
|
|
self,
|
|
command: str,
|
|
cwd: str | None = None,
|
|
env: dict[str, str] | None = None,
|
|
timeout: int | None = 30,
|
|
shell: bool = True,
|
|
background: bool = False,
|
|
) -> dict[str, Any]:
|
|
if not _is_safe_command(command):
|
|
raise PermissionError("Blocked unsafe shell command.")
|
|
|
|
def _run() -> dict[str, Any]:
|
|
run_env = os.environ.copy()
|
|
if env:
|
|
run_env.update({str(k): str(v) for k, v in env.items()})
|
|
working_dir = os.path.abspath(cwd) if cwd else get_astrbot_root()
|
|
if background:
|
|
# `command` is intentionally executed through the current shell so
|
|
# local computer-use behavior matches existing tool semantics.
|
|
# Safety relies on `_is_safe_command()` and the allowed-root checks.
|
|
proc = subprocess.Popen( # noqa: S602 # nosemgrep: python.lang.security.audit.dangerous-subprocess-use-audit
|
|
command,
|
|
shell=shell,
|
|
cwd=working_dir,
|
|
env=run_env,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
)
|
|
return {"pid": proc.pid, "stdout": "", "stderr": "", "exit_code": None}
|
|
# `command` is intentionally executed through the current shell so
|
|
# local computer-use behavior matches existing tool semantics.
|
|
# Safety relies on `_is_safe_command()` and the allowed-root checks.
|
|
result = subprocess.run( # noqa: S602 # nosemgrep: python.lang.security.audit.dangerous-subprocess-use-audit
|
|
command,
|
|
shell=shell,
|
|
cwd=working_dir,
|
|
env=run_env,
|
|
timeout=timeout,
|
|
capture_output=True,
|
|
)
|
|
return {
|
|
"stdout": _decode_shell_output(result.stdout),
|
|
"stderr": _decode_shell_output(result.stderr),
|
|
"exit_code": result.returncode,
|
|
}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
|
|
@dataclass
|
|
class LocalPythonComponent(PythonComponent):
|
|
async def exec(
|
|
self,
|
|
code: str,
|
|
kernel_id: str | None = None,
|
|
timeout: int = 30,
|
|
silent: bool = False,
|
|
) -> dict[str, Any]:
|
|
def _run() -> dict[str, Any]:
|
|
try:
|
|
result = subprocess.run(
|
|
[os.environ.get("PYTHON", sys.executable), "-c", code],
|
|
timeout=timeout,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
stdout = "" if silent else result.stdout
|
|
stderr = result.stderr if result.returncode != 0 else ""
|
|
return {
|
|
"data": {
|
|
"output": {"text": stdout, "images": []},
|
|
"error": stderr,
|
|
}
|
|
}
|
|
except subprocess.TimeoutExpired:
|
|
return {
|
|
"data": {
|
|
"output": {"text": "", "images": []},
|
|
"error": "Execution timed out.",
|
|
}
|
|
}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
|
|
@dataclass
|
|
class LocalFileSystemComponent(FileSystemComponent):
|
|
async def create_file(
|
|
self, path: str, content: str = "", mode: int = 0o644
|
|
) -> dict[str, Any]:
|
|
def _run() -> dict[str, Any]:
|
|
abs_path = os.path.abspath(path)
|
|
os.makedirs(os.path.dirname(abs_path), exist_ok=True)
|
|
with open(abs_path, "w", encoding="utf-8") as f:
|
|
f.write(content)
|
|
os.chmod(abs_path, mode)
|
|
return {"success": True, "path": abs_path}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
async def read_file(
|
|
self,
|
|
path: str,
|
|
encoding: str = "utf-8",
|
|
offset: int | None = None,
|
|
limit: int | None = None,
|
|
) -> dict[str, Any]:
|
|
def _run() -> dict[str, Any]:
|
|
abs_path = os.path.abspath(path)
|
|
detected_encoding = encoding
|
|
if encoding == "utf-8":
|
|
with open(abs_path, "rb") as f:
|
|
raw_sample = f.read(8192)
|
|
detected_encoding = detect_text_encoding(raw_sample) or encoding
|
|
return {
|
|
"success": True,
|
|
"content": read_local_text_range_sync(
|
|
abs_path,
|
|
encoding=detected_encoding,
|
|
offset=offset,
|
|
limit=limit,
|
|
),
|
|
}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
async def search_files(
|
|
self,
|
|
pattern: str,
|
|
path: str | None = None,
|
|
glob: str | None = None,
|
|
after_context: int | None = None,
|
|
before_context: int | None = None,
|
|
) -> dict[str, Any]:
|
|
def _run() -> dict[str, Any]:
|
|
results = search(
|
|
patterns=[pattern],
|
|
paths=[path] if path else None,
|
|
globs=[glob] if glob else None,
|
|
after_context=after_context,
|
|
before_context=before_context,
|
|
line_number=True,
|
|
)
|
|
return {"success": True, "content": _truncate_long_lines("".join(results))}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
async def edit_file(
|
|
self,
|
|
path: str,
|
|
old_string: str,
|
|
new_string: str,
|
|
replace_all: bool = False,
|
|
encoding: str = "utf-8",
|
|
) -> dict[str, Any]:
|
|
def _run() -> dict[str, Any]:
|
|
abs_path = os.path.abspath(path)
|
|
with open(abs_path, encoding=encoding) as f:
|
|
content = f.read()
|
|
occurrences = content.count(old_string)
|
|
if occurrences == 0:
|
|
return {
|
|
"success": False,
|
|
"error": "old string not found in file",
|
|
"replacements": 0,
|
|
}
|
|
if replace_all:
|
|
updated = content.replace(old_string, new_string)
|
|
replacements = occurrences
|
|
else:
|
|
updated = content.replace(old_string, new_string, 1)
|
|
replacements = 1
|
|
with open(abs_path, "w", encoding=encoding) as f:
|
|
f.write(updated)
|
|
return {
|
|
"success": True,
|
|
"path": abs_path,
|
|
"replacements": replacements,
|
|
}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
async def write_file(
|
|
self, path: str, content: str, mode: str = "w", encoding: str = "utf-8"
|
|
) -> dict[str, Any]:
|
|
def _run() -> dict[str, Any]:
|
|
abs_path = os.path.abspath(path)
|
|
os.makedirs(os.path.dirname(abs_path), exist_ok=True)
|
|
with open(abs_path, mode, encoding=encoding) as f:
|
|
f.write(content)
|
|
return {"success": True, "path": abs_path}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
async def delete_file(self, path: str) -> dict[str, Any]:
|
|
def _run() -> dict[str, Any]:
|
|
abs_path = os.path.abspath(path)
|
|
if os.path.isdir(abs_path):
|
|
shutil.rmtree(abs_path)
|
|
else:
|
|
os.remove(abs_path)
|
|
return {"success": True, "path": abs_path}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
async def list_dir(
|
|
self, path: str = ".", show_hidden: bool = False
|
|
) -> dict[str, Any]:
|
|
def _run() -> dict[str, Any]:
|
|
abs_path = os.path.abspath(path)
|
|
entries = os.listdir(abs_path)
|
|
if not show_hidden:
|
|
entries = [e for e in entries if not e.startswith(".")]
|
|
return {"success": True, "entries": entries}
|
|
|
|
return await asyncio.to_thread(_run)
|
|
|
|
|
|
class LocalBooter(ComputerBooter):
|
|
def __init__(self) -> None:
|
|
self._fs = LocalFileSystemComponent()
|
|
self._python = LocalPythonComponent()
|
|
self._shell = LocalShellComponent()
|
|
|
|
async def boot(self, session_id: str) -> None:
|
|
logger.info(f"Local computer booter initialized for session: {session_id}")
|
|
|
|
async def shutdown(self) -> None:
|
|
logger.info("Local computer booter shutdown complete.")
|
|
|
|
@property
|
|
def fs(self) -> FileSystemComponent:
|
|
return self._fs
|
|
|
|
@property
|
|
def python(self) -> PythonComponent:
|
|
return self._python
|
|
|
|
@property
|
|
def shell(self) -> ShellComponent:
|
|
return self._shell
|
|
|
|
async def upload_file(self, path: str, file_name: str) -> dict:
|
|
raise NotImplementedError(
|
|
"LocalBooter does not support upload_file operation. Use shell instead."
|
|
)
|
|
|
|
async def download_file(self, remote_path: str, local_path: str) -> None:
|
|
raise NotImplementedError(
|
|
"LocalBooter does not support download_file operation. Use shell instead."
|
|
)
|
|
|
|
async def available(self) -> bool:
|
|
return True
|