140 lines
4.8 KiB
Python
140 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
|
|
class ToolSandbox:
|
|
"""Applies simple working directory and timeout restrictions."""
|
|
|
|
def __init__(
|
|
self,
|
|
allowed_root: str | Path,
|
|
timeout_ms: int,
|
|
command_timeout_ms: int | None = None,
|
|
idle_timeout_ms: int | None = None,
|
|
) -> None:
|
|
self._allowed_root = Path(allowed_root).resolve()
|
|
self._timeout_seconds = max(timeout_ms / 1000, 0.001)
|
|
self._command_timeout_seconds = max((command_timeout_ms or timeout_ms) / 1000, 0.001)
|
|
self._idle_timeout_seconds = max((idle_timeout_ms or timeout_ms) / 1000, 0.001)
|
|
|
|
def ensure_path_allowed(self, path: str | Path) -> Path:
|
|
resolved = Path(path).expanduser().resolve()
|
|
# Permission-first model: path is allowed if it exists
|
|
# Permission service will handle write/shell restrictions
|
|
return resolved
|
|
|
|
def run_shell(
|
|
self,
|
|
command: str,
|
|
cwd: str | Path | None = None,
|
|
stdin_data: str | None = None,
|
|
output_callback: Callable[[str, str], None] | None = None,
|
|
) -> subprocess.CompletedProcess[str]:
|
|
working_directory = self.ensure_path_allowed(cwd or self._allowed_root)
|
|
env = {"PATH": os.environ.get("PATH", "")}
|
|
if output_callback is None:
|
|
return subprocess.run(
|
|
command,
|
|
shell=True,
|
|
cwd=str(working_directory),
|
|
env=env,
|
|
text=True,
|
|
capture_output=True,
|
|
input=stdin_data,
|
|
timeout=self._command_timeout_seconds,
|
|
check=False,
|
|
)
|
|
|
|
process = subprocess.Popen(
|
|
command,
|
|
shell=True,
|
|
cwd=str(working_directory),
|
|
env=env,
|
|
text=True,
|
|
stdin=subprocess.PIPE if stdin_data is not None else None,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
start_new_session=True,
|
|
)
|
|
|
|
stdout_chunks: list[str] = []
|
|
stderr_chunks: list[str] = []
|
|
output_lock = threading.Lock()
|
|
last_output_at = time.monotonic()
|
|
|
|
if stdin_data is not None and process.stdin is not None:
|
|
process.stdin.write(stdin_data)
|
|
process.stdin.close()
|
|
|
|
def read_stream(stream_name: str) -> None:
|
|
stream = process.stdout if stream_name == "stdout" else process.stderr
|
|
if stream is None:
|
|
return
|
|
chunks = stdout_chunks if stream_name == "stdout" else stderr_chunks
|
|
try:
|
|
for line in iter(stream.readline, ""):
|
|
if not line:
|
|
break
|
|
chunks.append(line)
|
|
nonlocal last_output_at
|
|
with output_lock:
|
|
last_output_at = time.monotonic()
|
|
output_callback(stream_name, line)
|
|
finally:
|
|
stream.close()
|
|
|
|
stdout_thread = threading.Thread(target=read_stream, args=("stdout",), daemon=True)
|
|
stderr_thread = threading.Thread(target=read_stream, args=("stderr",), daemon=True)
|
|
stdout_thread.start()
|
|
stderr_thread.start()
|
|
|
|
timed_out = False
|
|
timeout_reason: str | None = None
|
|
started_at = time.monotonic()
|
|
return_code: int | None = None
|
|
while return_code is None:
|
|
return_code = process.poll()
|
|
if return_code is not None:
|
|
break
|
|
|
|
now = time.monotonic()
|
|
with output_lock:
|
|
idle_for = now - last_output_at
|
|
if now - started_at > self._command_timeout_seconds:
|
|
timed_out = True
|
|
timeout_reason = f"Command timed out after {self._command_timeout_seconds:.0f}s"
|
|
break
|
|
if idle_for > self._idle_timeout_seconds:
|
|
timed_out = True
|
|
timeout_reason = f"Command produced no output for {self._idle_timeout_seconds:.0f}s"
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
if timed_out:
|
|
try:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
pass
|
|
except PermissionError:
|
|
process.kill()
|
|
return_code = process.wait()
|
|
timeout_message = f"{timeout_reason}\n"
|
|
stderr_chunks.append(timeout_message)
|
|
output_callback("stderr", timeout_message)
|
|
|
|
stdout_thread.join(timeout=1)
|
|
stderr_thread.join(timeout=1)
|
|
return subprocess.CompletedProcess(
|
|
args=command,
|
|
returncode=return_code if not timed_out else -9,
|
|
stdout="".join(stdout_chunks),
|
|
stderr="".join(stderr_chunks),
|
|
)
|