unsloth/studio/backend/core/inference/tools.py
Daniel Han aab371a068
studio: tighten sandbox blocklist precision (bash, hf upload, NOFILE) (#5487)
* studio: tighten sandbox blocklist precision (bash, hf upload, NOFILE)

Three precision fixes in core/inference/tools.py. Same security
boundary; fewer false positives that broke legitimate sandbox use.

bash blocklist:
The per-token loop introduced in #5375 fired on any blocklist word in
any token position, so the entirely benign `grep -r curl .`,
`echo source the data`, and `ls /usr/bin/curl` were rejected with
"blocked command 'curl'". The position-anchored regex already covers
real command-position invocations, including `;rm`, `&&wget`, `$(rm)`,
`<(rm)`, backticked subshells, and `/usr/bin/sudo`. The token loop is
re-scoped: it only fires when the previous shlex token is a shell
separator (or at start of line), so split-quoting obfuscations like
`r''m -rf /` are still caught (shlex collapses them to a single
command-position token) while argument-position blocklist words pass
through. Trailing meta-chars glued to a shlex token (`rm;`) are
stripped before basename matching.

hf upload AST gate:
`_method_call_is_hf_upload` previously matched any method named
`upload_file` / `upload_folder` / `upload_large_folder` / `create_commit`
on any receiver, so paramiko.SFTPClient.upload_file, boto3.create_commit,
and similar non-HF SDK methods were rejected. The fallback now requires
an `import huggingface_hub` / `import hf_api` / `from huggingface_hub
import ...` somewhere in the same module. Fully-qualified
huggingface_hub.upload_file(...) calls are unchanged.

NOFILE env knob:
`RLIMIT_NOFILE = (1024, 1024)` was the only sandbox rlimit without an
env override. 1024 is below Linux's typical soft default and below
what multi-shard safetensors mmap chains need on Llama-3 70B-class
loads. Default is now 16384 with UNSLOTH_STUDIO_SANDBOX_NOFILE, parity
with the other rlimits.

15 new bash-blocklist-position tests pin both the false-positive
fixes and the still-blocked invariants (semicolon, &&, subshell,
backtick, split-quote, /usr/bin/ prefix, nested bash -c).
4 new hf-upload-import-gate tests pin both the false-positive
allowances and that HF-imported uses are still blocked.
1 new pin asserts the NOFILE env var is wired.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: cover command wrappers, find -exec, dynamic HF imports, NOFILE clamp

Reviewer follow-ups to the sandbox blocklist precision change.

Command-position scanner missed Bash command-prefix wrappers and inline
shell assignments. shlex tokenised `env curl`, `time curl`, `nohup rm`,
`FOO=bar curl`, `sudo rm`, etc. with the prefix at command position and
the real command at argument position, so the position-anchored check
returned set() while pre-PR's per-token scan caught them. Likewise the
position-anchored regex requires `^` or a shell separator before the
command, so `env curl` slipped through.

Reworked the scanner to track an expect_command flag plus a
prefix_pending flag:
  - assignments (FOO=bar) keep expect_command=True for the next token,
  - flags ('-oL', '--') keep it intact while prefix_pending is set,
  - numeric duration args ('timeout 1 cmd') skip without breaking
    expect_command,
  - known wrappers (env, command, builtin, exec, time, nohup, nice,
    setsid, stdbuf, timeout, ionice, chroot, sudo, doas, su, xargs)
    set prefix_pending so the wrapper's command is still checked,
  - shell separators now include `{`, `}`, `)`, `then`, `do`,
    `else`, `elif` so brace groups and if/then/while/do bodies are
    recognised as command positions.

Also lex with `shlex.shlex(punctuation_chars=";&|()`")` so split-quote
forms like `echo done; r''m -rf /tmp/x` and `echo done;r''m` tokenise
as `[..., ';', 'rm', ...]` and the command position check fires.

Added a small `find -exec CMD ... ;` / `-execdir CMD ... ;` pass so
`find . -exec rm -f {} +` and friends are caught even though the
direct token is at argument position to `find`.

Dynamic Hugging Face imports were treated as no-HF-in-scope. The
upload-method gate now also resolves `__import__('huggingface_hub')`,
`importlib.import_module('huggingface_hub')`, and bare
`import_module('huggingface_hub')` (via `from importlib import
import_module`) as HF imports, so HfApi().upload_file via dynamic
import is still blocked.

RLIMIT_NOFILE: setrlimit(NOFILE, (16384, 16384)) silently failed if
the parent's hard cap is below the requested value; the broad
except swallowed the OSError and left the sandbox at the parent's
default. Clamp the requested value to the inherited hard limit
before calling setrlimit.

Test cleanup: the existing test_cat_with_word_source_allowed had
`assert ... or True` so it could not fail; rewrote it to assert the
actual return value plus the two membership checks. Added
parametrised coverage for shell prefix wrappers, find -exec / xargs,
brace groups, if/then, while/do, split-quote command-name forms, and
dynamic HF import upload patterns.

Test:
  - python -m pytest studio/backend/tests/test_sandbox_tools.py -q
    -> 90 passed (was 67 before this commit)
  - full studio/backend/tests/ minus llama_cpp_load_progress_live and
    GPU CUDA_VISIBLE_DEVICES tests (pre-existing isolation flake)
    -> 1063 passed

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: catch bare-name HF upload calls in AST gate

`from huggingface_hub import upload_file; upload_file(...)` is a
canonical HF call shape that the previous Attribute-only check missed:
the bare-name call lands as ast.Name (not ast.Attribute), so the
fuzzy gate skipped it.

Extend _method_call_is_hf_upload to also match ast.Name when HF is in
scope. Same import-gating discipline as the Attribute branch, so
paramiko/boto3 and locally-defined `def upload_file(...)` helpers
without HF imports still pass.

Pins: 4 new TestHfUploadImportGate cases (upload_file/folder/create_commit
bare-name imports blocked; local upload_file without HF import allowed).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: scope HF uploads to sandbox-local literals; block env / token leaks

The previous gate dropped every HF upload call. Two refinements make it
precise enough to allow legitimate sandbox->HF uploads while still
catching credential / file exfil:

- path_or_fileobj / folder_path / create_commit operation paths must be
  sandbox-local relative-path literals (no '/', '~', drive letter, or
  '..' segments). Variable / dynamic paths are rejected.

- Any positional or keyword argument that statically resolves to
  os.environ / os.environ.get / os.getenv / bare getenv / subprocess
  shape readers is rejected (env-var exfil).

- token / hf_token / api_token / api_key / auth_token / access_token /
  password / secret kwargs are always rejected; sandbox env strips all
  parent credentials by construction, so any value here is hard-coded
  or lifted.

Recursive subtree walk in _reads_env_or_secret catches wrapper shapes
(str(os.environ), json.dumps(os.environ.items()), etc.).

Add TestSandboxEnvIsolation: pin that _build_safe_env builds the env
from a whitelist, not by stripping. Cover Linux/macOS/WSL/Windows
secret shapes. The whitelist is PATH / HOME / TMPDIR / LANG / TERM /
PYTHONIOENCODING (+ VIRTUAL_ENV / SystemRoot when applicable); HOME
points at the sandbox workdir, so HF / wandb / aws SDKs cannot reach
the operator's ~/.cache credentials.

Test classes added:
- TestHfUploadSandboxLocalPaths (relative literals allowed; absolute,
  drive-letter, '~', '..', mid-path traversal, dynamic vars, and
  open() of unsafe paths blocked, including create_commit recursion).
- TestHfUploadEnvAndSecretLeakBlock (os.environ subscript/get/getenv,
  bare getenv, subprocess.check_output, str(os.environ), token=,
  hf_token=, api_key=, and create_commit operations referencing env).
- TestSandboxEnvIsolation (no parent secret leaks into sandbox env).

131 tests in test_sandbox_tools.py pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-18 00:01:17 -07:00

2026 lines
73 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Tool definitions and executors for LLM tool calling.
Supports web search (DuckDuckGo), Python code execution, and terminal commands.
"""
import ast
import http.client
import os
import signal
os.environ["UNSLOTH_IS_PRESENT"] = "1"
import random
import re
import shlex
import ssl
import subprocess
import sys
import tempfile
import threading
import urllib.request
from loggers import get_logger
logger = get_logger(__name__)
_EXEC_TIMEOUT = 300 # 5 minutes
# Pre-import modules used in _sandbox_preexec at module level so that
# the preexec_fn closure does not trigger the import machinery in the
# forked child (which can deadlock in multi-threaded servers).
_libc = None
if sys.platform == "linux":
try:
import ctypes
import ctypes.util
_libc_name = ctypes.util.find_library("c")
if _libc_name:
_libc = ctypes.CDLL(_libc_name, use_errno = True)
except (OSError, AttributeError):
pass
_resource = None
if sys.platform != "win32":
try:
import resource as _resource
except ImportError:
pass
# Strict raster-image allowlist for sandbox file serving.
# No .svg (XSS risk via embedded scripts), no .html, no .pdf.
_IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"})
_MAX_OUTPUT_CHARS = 8000 # truncate long output
_BLOCKED_COMMANDS_COMMON = frozenset(
{
"rm",
"dd",
"chmod",
"chown",
"mkfs",
"mount",
"umount",
"fdisk",
"sudo",
"su",
"doas",
"pkexec",
"shutdown",
"reboot",
"halt",
"poweroff",
"kill",
"killall",
"pkill",
"passwd",
"curl",
"wget",
"nc",
"ncat",
"netcat",
"socat",
"ssh",
"scp",
"sftp",
"rsync",
"eval",
"source",
}
)
_BLOCKED_COMMANDS_WIN = frozenset(
{
"rmdir",
"takeown",
"icacls",
"runas",
"powershell",
"pwsh",
}
)
_BLOCKED_COMMANDS = (
_BLOCKED_COMMANDS_COMMON | _BLOCKED_COMMANDS_WIN
if sys.platform == "win32"
else _BLOCKED_COMMANDS_COMMON
)
_SHELL_SEPARATORS = frozenset(
{";", "&&", "||", "|", "&", "\n", "(", ")", "`", "{", "}"}
)
# Bash keywords that introduce a new command position (then $cmd, do $cmd, etc.).
_SHELL_KEYWORDS_AS_SEP = frozenset({"then", "do", "else", "elif"})
# Wrappers whose next non-flag argument is itself the command Bash will exec.
_COMMAND_PREFIXES = frozenset(
{
"env",
"command",
"builtin",
"exec",
"time",
"nohup",
"nice",
"setsid",
"stdbuf",
"timeout",
"ionice",
"chroot",
"sudo",
"doas",
"su",
"xargs",
}
)
_ASSIGNMENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=")
_FIND_EXEC_FLAGS = frozenset({"-exec", "-execdir", "-ok", "-okdir"})
def _find_blocked_commands(command: str) -> set[str]:
"""Detect blocked commands at shell command position only.
A token is at command position if it is the first token, or if the
preceding token is a shell separator / brace-group opener / keyword
that starts a new command (`then`, `do`, etc.), or a command-prefix
wrapper like `env` / `time` / `xargs` (the next token is the real
command). Tokens in argument position (`grep -r curl .`,
`echo source the data`, `ls /usr/bin/curl`) are passed through.
Also scans `find ... -exec CMD` and recurses into bash -c / cmd /c.
"""
blocked: set[str] = set()
# shlex with punctuation_chars splits `;`, `&&`, `||`, `|`, `(`, `)`, `` ` ``
# off as their own tokens so we can detect command position even when a
# caller writes `echo done; rm -rf x` (no whitespace) or quote-splits the
# command name itself (`r''m` collapses to a single token `rm` at command
# position after the `;` separator).
try:
if sys.platform == "win32":
tokens = shlex.split(command, posix = False)
else:
lexer = shlex.shlex(command, posix = True, punctuation_chars = ";&|()`")
lexer.whitespace_split = True
tokens = list(lexer)
except ValueError:
tokens = command.split()
def _token_basename(tok: str) -> str:
# shlex may glue trailing meta-chars onto a token (`rm;`); strip them
# so the basename match still hits `rm`. Leading shell-state chars
# likewise.
tok = tok.strip(";&|()`{}")
base = os.path.basename(tok).lower()
stem, ext = os.path.splitext(base)
if ext in {".exe", ".com", ".bat", ".cmd"}:
base = stem
return base
expect_command = True # start of string is a command position
prefix_pending = False # last command-position token was env/time/timeout/xargs/...
for token in tokens:
if token in _SHELL_SEPARATORS or token in _SHELL_KEYWORDS_AS_SEP:
expect_command = True
prefix_pending = False
continue
if token.startswith("-"):
# Flags belong to the active command. While a wrapper prefix is
# waiting for its command (`stdbuf -oL cmd`, `xargs -- cmd`),
# keep expect_command intact.
if not prefix_pending:
expect_command = False
continue
if not expect_command:
continue
# FOO=bar prefix: assignment list, next non-assignment token is the command.
if _ASSIGNMENT_RE.match(token):
continue
# `timeout 1 cmd` / `nice -n 5 cmd` style numeric wrapper arg.
if prefix_pending and token.lstrip("-").isdigit():
continue
base = _token_basename(token)
if base in _BLOCKED_COMMANDS:
blocked.add(base)
# Wrappers (`env` / `time` / `xargs` / `sudo`) consume one command; the
# next non-flag, non-numeric token is the real command. `sudo` is
# already in _BLOCKED_COMMANDS, so it's flagged AND we keep walking.
if base in _COMMAND_PREFIXES:
prefix_pending = True
continue
expect_command = False
prefix_pending = False
# `find ... -exec CMD ... ;` and `-execdir CMD ... ;` invoke CMD directly.
for i, tok in enumerate(tokens):
if tok in _FIND_EXEC_FLAGS and i + 1 < len(tokens):
base = _token_basename(tokens[i + 1])
if base in _BLOCKED_COMMANDS:
blocked.add(base)
# Regex: blocked words at shell command boundaries that shlex won't see,
# e.g. inside an unquoted $(rm -rf), <(rm), backtick chain, or appended to
# a separator with no whitespace ("foo;rm"). Anchored to command-position
# delimiters; does not match in argument position.
lowered = command.lower()
if _BLOCKED_COMMANDS:
words_alt = "|".join(re.escape(w) for w in sorted(_BLOCKED_COMMANDS))
pattern = (
rf"(?:^|[;&|`\n(]\s*|[$]\(\s*|<\(\s*)"
rf"(?:[\w./\\-]*/|[a-zA-Z]:[/\\][\w./\\-]*)?"
rf"({words_alt})(?:\.(?:exe|com|bat|cmd))?\b"
)
blocked.update(re.findall(pattern, lowered))
# Nested shell invocations (bash -c 'sudo whoami',
# bash -lc '...', bash --login -c '...', cmd /c '...').
# When a -c or /c flag is found, look backwards for a shell name
# (skipping intermediate flags like --login, -l, -x) and recursively
# scan the nested command string.
_SHELLS = {"bash", "sh", "zsh", "dash", "ksh", "csh", "tcsh", "fish"}
_SHELLS_WIN = {"cmd", "cmd.exe"}
for i, token in enumerate(tokens):
tok_lower = token.lower()
# Match -c exactly, or combined flags ending in c (e.g. -lc, -xc)
is_unix_c = tok_lower == "-c" or (
tok_lower.startswith("-")
and tok_lower.endswith("c")
and not tok_lower.startswith("--")
)
is_win_c = tok_lower == "/c"
if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens):
continue
# Look backwards past any flags to find the shell binary.
# On Unix, flags start with - (skip those). On Windows, flags
# start with / but so do absolute paths, so only skip short
# single-char /X flags (not /bin/bash style paths).
for j in range(i - 1, -1, -1):
prev = tokens[j]
if prev.startswith("-"):
continue # skip Unix flags like --login, -l
if is_win_c and prev.startswith("/") and len(prev) <= 3:
continue # skip Windows flags like /s, /q (not /bin/bash)
prev_base = os.path.basename(prev).lower()
if is_unix_c and prev_base in _SHELLS:
blocked |= _find_blocked_commands(tokens[i + 1])
elif is_win_c and prev_base in _SHELLS_WIN:
blocked |= _find_blocked_commands(tokens[i + 1])
break # stop at first non-flag token
return blocked
def _build_safe_env(workdir: str) -> dict[str, str]:
"""Build a minimal, credential-free environment for sandboxed subprocesses.
Whitelist-built from scratch -- the parent process env is NOT inherited.
Only PATH / HOME / TMPDIR / LANG / TERM / PYTHONIOENCODING (+ VIRTUAL_ENV
or Windows SystemRoot when applicable) reach the child. HF_TOKEN,
WANDB_API_KEY, AWS_*, GH_TOKEN, OPENAI_API_KEY, LD_PRELOAD, DYLD_*, and
every other parent var are absent by construction. HOME points at the
sandbox workdir so HF / wandb / aws SDKs cannot read cached credentials
from the operator's real ~/.
"""
# Start with the directory containing the running Python interpreter
# so that subprocess calls to 'python', 'pip', etc. resolve to the
# same environment the Studio server is running in.
exe_dir = os.path.dirname(sys.executable)
path_entries = [exe_dir] if exe_dir else []
# If a virtualenv is active, include its bin/Scripts directory.
venv = os.environ.get("VIRTUAL_ENV")
if venv:
venv_bin = os.path.join(venv, "Scripts" if sys.platform == "win32" else "bin")
if venv_bin not in path_entries:
path_entries.append(venv_bin)
if sys.platform == "win32":
sysroot = os.environ.get("SystemRoot", r"C:\Windows")
path_entries.extend([os.path.join(sysroot, "System32"), sysroot])
else:
path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"])
# Deduplicate while preserving order
deduped = list(dict.fromkeys(p for p in path_entries if p))
env = {
"PATH": os.pathsep.join(deduped),
"HOME": workdir,
"TMPDIR": workdir,
"LANG": os.environ.get("LANG", "C.UTF-8"),
"TERM": "dumb",
"PYTHONIOENCODING": "utf-8",
}
if venv:
env["VIRTUAL_ENV"] = venv
# Windows needs SystemRoot for Python/subprocess to work
if sys.platform == "win32":
env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows")
return env
def _sandbox_preexec():
"""Best-effort sandbox setup for sandboxed subprocesses.
Modules are resolved at import time so the forked child runs no imports.
"""
try:
os.setsid()
except OSError:
pass
try:
os.umask(0o077)
except OSError:
pass
if _libc is not None:
try:
_libc.prctl(38, 1, 0, 0, 0) # PR_SET_NO_NEW_PRIVS
except (OSError, AttributeError):
pass
try:
_libc.prctl(1, 9, 0, 0, 0) # PR_SET_PDEATHSIG = SIGKILL
except (OSError, AttributeError):
pass
# CLONE_NEWNET intentionally not applied: where userns is enabled it
# blocks all egress, including allowlisted hosts. Network policy is
# enforced by the AST host check and the bash blocklist.
if _resource is not None:
# RLIMIT_NPROC is per-real-UID, so the cap is well above normal usage.
try:
nproc = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_NPROC", "10000"))
_resource.setrlimit(_resource.RLIMIT_NPROC, (nproc, nproc))
except (ValueError, OSError, AttributeError):
pass
try:
_resource.setrlimit(
_resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024)
)
except (ValueError, OSError):
pass
try:
as_bytes = (
int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_AS_GB", "8"))
* 1024
* 1024
* 1024
)
_resource.setrlimit(_resource.RLIMIT_AS, (as_bytes, as_bytes))
except (ValueError, OSError, AttributeError):
pass
try:
cpu_s = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_CPU_S", "600"))
_resource.setrlimit(_resource.RLIMIT_CPU, (cpu_s, cpu_s))
except (ValueError, OSError, AttributeError):
pass
try:
# Default high enough for multi-shard safetensors mmaps + Python's
# own handle count; tunable via env for installs that hit the cap.
# Clamp to the inherited hard limit so setrlimit doesn't ValueError
# on machines where the parent's hard cap is below the requested
# value (would otherwise leave NOFILE at the parent's default).
nofile = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_NOFILE", "16384"))
_soft_cur, hard_cur = _resource.getrlimit(_resource.RLIMIT_NOFILE)
target = (
nofile if hard_cur == _resource.RLIM_INFINITY else min(nofile, hard_cur)
)
_resource.setrlimit(_resource.RLIMIT_NOFILE, (target, target))
except (ValueError, OSError, AttributeError):
pass
def _get_shell_cmd(command: str) -> list[str]:
"""Return the platform-appropriate shell invocation for a command string."""
if sys.platform == "win32":
return ["cmd", "/c", command]
return ["bash", "-c", command]
# Per-session working directories so each chat thread gets its own sandbox.
# Falls back to a shared ~/studio_sandbox/_default for API callers without a
# session_id.
_workdirs: dict[str, str] = {}
# Non-matching session_ids collapse to ``_invalid`` to block cross-session escapes.
_SESSION_ID_RE = re.compile(r"\A[A-Za-z0-9_\-]{1,64}\Z")
def _get_workdir(session_id: str | None = None) -> str:
"""Return a per-session sandbox dir at mode 0o700."""
global _workdirs
key = session_id or "_default"
if key not in _workdirs or not os.path.isdir(_workdirs[key]):
home = os.path.expanduser("~")
sandbox_root = os.path.join(home, "studio_sandbox")
if session_id and _SESSION_ID_RE.match(session_id):
workdir = os.path.join(sandbox_root, session_id)
if not os.path.realpath(workdir).startswith(
os.path.realpath(sandbox_root) + os.sep
):
workdir = os.path.join(sandbox_root, "_invalid")
elif session_id:
workdir = os.path.join(sandbox_root, "_invalid")
else:
workdir = os.path.join(sandbox_root, "_default")
os.makedirs(workdir, exist_ok = True)
try:
os.chmod(sandbox_root, 0o700)
except OSError:
pass
try:
os.chmod(workdir, 0o700)
except OSError:
pass
_workdirs[key] = workdir
return _workdirs[key]
WEB_SEARCH_TOOL = {
"type": "function",
"function": {
"name": "web_search",
"description": (
"Search the web and fetch page content. Returns snippets for all results. "
"Use the url parameter to fetch full page text from a specific URL."
),
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query",
},
"url": {
"type": "string",
"description": "A URL to fetch full page content from (instead of searching). Use this to read a page found in search results.",
},
},
"required": [],
},
},
}
PYTHON_TOOL = {
"type": "function",
"function": {
"name": "python",
"description": "Execute Python code in a sandbox and return stdout/stderr.",
"parameters": {
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "The Python code to run",
}
},
"required": ["code"],
},
},
}
TERMINAL_TOOL = {
"type": "function",
"function": {
"name": "terminal",
"description": "Execute a terminal command and return stdout/stderr.",
"parameters": {
"type": "object",
"properties": {
"command": {
"type": "string",
"description": "The command to run",
}
},
"required": ["command"],
},
},
}
ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL]
_TIMEOUT_UNSET = object()
def execute_tool(
name: str,
arguments: dict,
cancel_event = None,
timeout: int | None = _TIMEOUT_UNSET,
session_id: str | None = None,
) -> str:
"""Execute a tool by name with the given arguments. Returns result as a string.
``timeout``: int sets per-call limit in seconds, ``None`` means no limit,
unset (default) uses ``_EXEC_TIMEOUT`` (300 s).
``session_id``: optional thread/session ID for per-conversation sandbox isolation.
"""
logger.info(
f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}"
)
effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
if name == "web_search":
return _web_search(
arguments.get("query", ""),
url = arguments.get("url"),
timeout = effective_timeout,
)
if name == "python":
return _python_exec(
arguments.get("code", ""), cancel_event, effective_timeout, session_id
)
if name == "terminal":
return _bash_exec(
arguments.get("command", ""), cancel_event, effective_timeout, session_id
)
return f"Unknown tool: {name}"
_MAX_PAGE_CHARS = 16000 # limit fetched page text (after HTML-to-MD conversion)
# Raw download cap. Must be larger than _MAX_PAGE_CHARS because SSR pages
# embed large <head> sections (CSS, JS, SVGs) that are stripped during
# HTML-to-Markdown conversion. 512 KB is enough to reach article content
# on GitBook / Next.js / Docusaurus pages whose <head> alone can be 200 KB.
_MAX_FETCH_BYTES = 512 * 1024
_USER_AGENTS = (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:133.0) Gecko/20100101 Firefox/133.0",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:133.0) Gecko/20100101 Firefox/133.0",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/18.2 Safari/605.1.15",
)
_tls_ctx = ssl.create_default_context()
class _NoRedirect(urllib.request.HTTPRedirectHandler):
def redirect_request(self, req, fp, code, msg, headers, newurl):
return None
class _PinnedHTTPSConnection(http.client.HTTPSConnection):
"""HTTPS connection that connects to a pinned IP but uses a different
hostname for SNI and certificate verification.
The SSRF IP-pinning rewrites URLs to raw IPs. A normal HTTPSConnection
would then send no SNI and verify the cert against the IP, both of which
fail. This subclass splits the two concerns: TCP connects to the pinned
IP (``host`` parameter) while TLS uses ``sni_hostname`` for the
ClientHello and cert check.
"""
def __init__(self, host: str, *, sni_hostname: str, **kwargs):
super().__init__(host, **kwargs)
self._sni_hostname = sni_hostname
def connect(self):
# TCP connect to the pinned IP stored in self.host (+ tunnel if
# a proxy is configured via set_tunnel, though we do not use one).
http.client.HTTPConnection.connect(self)
# TLS handshake with the real hostname for SNI + cert verification.
self.sock = self._context.wrap_socket(
self.sock,
server_hostname = self._sni_hostname,
)
class _SNIHTTPSHandler(urllib.request.HTTPSHandler):
"""HTTPS handler that sends the correct SNI hostname during TLS handshake.
The SSRF IP-pinning rewrites URLs to raw IPs, which breaks SNI and cert
verification. This handler returns a ``_PinnedHTTPSConnection`` that
connects to the pinned IP but verifies TLS against the original hostname.
"""
def __init__(self, hostname: str):
super().__init__(context = _tls_ctx)
self._sni_hostname = hostname
def https_open(self, req):
return self.do_open(self._sni_connection, req)
def _sni_connection(self, host, **kwargs):
kwargs["context"] = _tls_ctx
return _PinnedHTTPSConnection(host, sni_hostname = self._sni_hostname, **kwargs)
def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str]:
"""Resolve *hostname*, reject non-public IPs, return a pinned IP string.
Returns ``(ok, reason_or_empty, resolved_ip)``. The caller should
connect to *resolved_ip* (with a ``Host`` header) to prevent DNS
rebinding between validation and the actual fetch.
"""
import ipaddress
import socket
try:
infos = socket.getaddrinfo(hostname, port, type = socket.SOCK_STREAM)
except OSError as e:
return False, f"Failed to resolve host: {e}", ""
if not infos:
return False, f"Failed to resolve host: no addresses for {hostname!r}", ""
for *_, sockaddr in infos:
ip = ipaddress.ip_address(sockaddr[0])
if (
ip.is_private
or ip.is_loopback
or ip.is_link_local
or ip.is_multicast
or ip.is_reserved
or ip.is_unspecified
):
return False, f"Blocked: refusing to fetch non-public address {ip}.", ""
# Return the first resolved address for pinning
first_ip = infos[0][4][0]
return True, "", first_ip
def _fetch_page_text(
url: str, max_chars: int = _MAX_PAGE_CHARS, timeout: int = 30
) -> str:
"""Fetch a URL and return plain text content (HTML tags stripped).
Blocks private/loopback/link-local targets (SSRF protection) and caps
the download size to avoid unbounded memory usage.
"""
from urllib.parse import urlparse
parsed = urlparse(url)
if parsed.scheme not in ("http", "https"):
return f"Blocked: only http/https URLs are allowed (got {parsed.scheme!r})."
if not parsed.hostname:
return "Blocked: URL is missing a hostname."
port = parsed.port or (443 if parsed.scheme == "https" else 80)
ok, reason, pinned_ip = _validate_and_resolve_host(parsed.hostname, port)
if not ok:
return reason
try:
from urllib.error import HTTPError as _HTTPError
from urllib.parse import urljoin, urlunparse
max_bytes = _MAX_FETCH_BYTES
current_url = url
current_host = parsed.hostname
ua = random.choice(_USER_AGENTS)
for _hop in range(5):
# Pin to the validated IP to prevent DNS rebinding.
# Rewrite the URL to use the IP and set the Host header.
cp = urlparse(current_url)
# Bracket IPv6 addresses so the netloc is valid in a URL.
ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip
ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str
pinned_url = urlunparse(cp._replace(netloc = ip_netloc))
opener = urllib.request.build_opener(
_NoRedirect,
_SNIHTTPSHandler(current_host),
)
req = urllib.request.Request(
pinned_url,
headers = {
"User-Agent": ua,
"Host": current_host,
},
)
try:
resp = opener.open(req, timeout = timeout)
except _HTTPError as e:
if e.code not in (301, 302, 303, 307, 308):
return (
f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}"
)
location = e.headers.get("Location")
if not location:
return "Failed to fetch URL: redirect missing Location header."
current_url = urljoin(current_url, location)
rp = urlparse(current_url)
if rp.scheme not in ("http", "https") or not rp.hostname:
return "Blocked: redirect target is not a valid http/https URL."
rp_port = rp.port or (443 if rp.scheme == "https" else 80)
ok2, reason2, pinned_ip = _validate_and_resolve_host(
rp.hostname,
rp_port,
)
if not ok2:
return reason2
current_host = rp.hostname
continue
# Success -- read capped body
raw_bytes = resp.read(max_bytes)
break
else:
return "Failed to fetch URL: too many redirects."
charset = resp.headers.get_content_charset() or "utf-8"
raw_html = raw_bytes.decode(charset, errors = "replace")
except _HTTPError as e:
return f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}"
except Exception as e:
return f"Failed to fetch URL: {e}"
# Convert HTML to Markdown using the builtin converter (no external deps)
from ._html_to_md import html_to_markdown
text = html_to_markdown(raw_html)
if not text:
return "(page returned no readable text)"
if len(text) > max_chars:
text = text[:max_chars] + f"\n\n... (truncated, {len(text)} chars total)"
return text
def _web_search(
query: str,
max_results: int = 5,
timeout: int = _EXEC_TIMEOUT,
url: str | None = None,
) -> str:
"""Search the web using DuckDuckGo and return formatted results.
If ``url`` is provided, fetches that page directly instead of searching.
"""
# Direct URL fetch mode
if url and url.strip():
fetch_timeout = 60 if timeout is None else min(timeout, 60)
return _fetch_page_text(url.strip(), timeout = fetch_timeout)
if not query or not query.strip():
return "No query provided."
try:
from ddgs import DDGS
results = DDGS(timeout = timeout).text(query, max_results = max_results)
if not results:
return "No results found."
parts = []
for r in results:
parts.append(
f"Title: {r.get('title', '')}\n"
f"URL: {r.get('href', '')}\n"
f"Snippet: {r.get('body', '')}"
)
text = "\n\n---\n\n".join(parts)
text += (
"\n\n---\n\nIMPORTANT: These are only short snippets. "
"To get the full page content, call web_search with "
'the url parameter (e.g. {"url": "<URL>"}).'
)
return text
except Exception as e:
return f"Search failed: {e}"
def _check_signal_escape_patterns(code: str):
"""
Check if code contains patterns that could escape signal-based timeouts.
Vendored from unsloth_zoo.rl_environments to avoid importing unsloth_zoo
(which requires GPU drivers and fails on Mac/Apple Silicon).
Returns (safe: bool, details: dict)
"""
try:
tree = ast.parse(code)
except SyntaxError as e:
return False, {
"error": f"SyntaxError: {e}",
"signal_tampering": [],
"exception_catching": [],
"warnings": [],
}
signal_tampering = []
exception_catching = []
shell_escapes = []
warnings = []
def _ast_name_matches(node, names):
if isinstance(node, ast.Name):
return node.id in names
elif isinstance(node, ast.Attribute):
full_name = []
current = node
while isinstance(current, ast.Attribute):
full_name.append(current.attr)
current = current.value
if isinstance(current, ast.Name):
full_name.append(current.id)
full_name = ".".join(reversed(full_name))
return full_name in names
return False
# Dangerous os/subprocess functions that can execute shell commands
_SHELL_EXEC_FUNCS = frozenset(
{
"os.system",
"os.popen",
"os.popen2",
"os.popen3",
"os.popen4",
"os.execl",
"os.execle",
"os.execlp",
"os.execlpe",
"os.execv",
"os.execve",
"os.execvp",
"os.execvpe",
"os.spawnl",
"os.spawnle",
"os.spawnlp",
"os.spawnlpe",
"os.spawnv",
"os.spawnve",
"os.spawnvp",
"os.spawnvpe",
"os.posix_spawn",
"os.posix_spawnp",
"subprocess.run",
"subprocess.call",
"subprocess.check_call",
"subprocess.check_output",
"subprocess.Popen",
"subprocess.getoutput",
"subprocess.getstatusoutput",
}
)
def _extract_string_from_node(node):
"""Extract a plain string value from an AST node, if it is a constant."""
if isinstance(node, ast.Constant) and isinstance(node.value, str):
return node.value
return None
def _extract_strings_from_list(node):
"""Extract string elements from an AST List or Tuple node."""
if isinstance(node, (ast.List, ast.Tuple)):
parts = []
for elt in node.elts:
s = _extract_string_from_node(elt)
if s is not None:
parts.append(s)
return parts
return []
# Keyword argument names that carry command content (as opposed to
# control flags like check=True, text=True, capture_output=True).
_CMD_KWARGS = frozenset({"args", "command", "executable", "path", "file"})
def _check_args_for_blocked(args_nodes):
"""Check if any call arguments contain blocked commands."""
found = set()
for arg in args_nodes:
s = _extract_string_from_node(arg)
if s is not None:
found |= _find_blocked_commands(s)
strs = _extract_strings_from_list(arg)
for s in strs:
found |= _find_blocked_commands(s)
return found
class SignalEscapeVisitor(ast.NodeVisitor):
def __init__(self):
self.imports_signal = False
self.signal_aliases = {"signal"}
self.os_aliases = {"os"}
self.subprocess_aliases = {"subprocess"}
# Maps bare function names to their fully-qualified form
# for from-import tracking (e.g. "system" -> "os.system")
self.shell_exec_aliases: dict[str, str] = {}
self.loop_depth = 0
def visit_Import(self, node):
for alias in node.names:
if alias.name == "signal":
self.imports_signal = True
if alias.asname:
self.signal_aliases.add(alias.asname)
elif alias.name == "os":
self.os_aliases.add(alias.asname or "os")
elif alias.name == "subprocess":
self.subprocess_aliases.add(alias.asname or "subprocess")
self.generic_visit(node)
def visit_ImportFrom(self, node):
if node.module == "signal":
self.imports_signal = True
for alias in node.names:
if alias.name in (
"signal",
"SIGALRM",
"SIG_IGN",
"setitimer",
"ITIMER_REAL",
"pthread_sigmask",
"SIG_BLOCK",
"alarm",
):
self.signal_aliases.add(alias.asname or alias.name)
elif node.module in ("os", "subprocess"):
if node.module == "os":
self.os_aliases.add("os")
else:
self.subprocess_aliases.add("subprocess")
# Track from-imports of dangerous functions
for alias in node.names:
fq = f"{node.module}.{alias.name}"
if fq in _SHELL_EXEC_FUNCS:
self.shell_exec_aliases[alias.asname or alias.name] = fq
self.generic_visit(node)
def visit_While(self, node):
self.loop_depth += 1
self.generic_visit(node)
self.loop_depth -= 1
def visit_For(self, node):
self.loop_depth += 1
self.generic_visit(node)
self.loop_depth -= 1
def visit_Call(self, node):
func = node.func
func_name = None
if isinstance(func, ast.Attribute):
if isinstance(func.value, ast.Name):
if func.value.id in self.signal_aliases:
func_name = f"signal.{func.attr}"
elif isinstance(func, ast.Name):
if func.id in ("signal", "setitimer", "alarm", "pthread_sigmask"):
func_name = func.id
if func_name:
if func_name in ("signal.signal", "signal"):
if len(node.args) >= 1:
if _ast_name_matches(
node.args[0], ("SIGALRM", "signal.SIGALRM")
):
signal_tampering.append(
{
"type": "signal_handler_override",
"line": node.lineno,
"description": "Overrides SIGALRM handler",
}
)
elif func_name in ("signal.setitimer", "setitimer"):
if len(node.args) >= 1:
if _ast_name_matches(
node.args[0], ("ITIMER_REAL", "signal.ITIMER_REAL")
):
signal_tampering.append(
{
"type": "timer_manipulation",
"line": node.lineno,
"description": "Manipulates ITIMER_REAL timer",
}
)
elif func_name in ("signal.alarm", "alarm"):
signal_tampering.append(
{
"type": "alarm_manipulation",
"line": node.lineno,
"description": "Manipulates alarm timer",
}
)
elif func_name in ("signal.pthread_sigmask", "pthread_sigmask"):
signal_tampering.append(
{
"type": "signal_mask",
"line": node.lineno,
"description": "Modifies signal mask (may block SIGALRM)",
}
)
# --- Shell escape detection ---
# Resolve the fully qualified function name for os.*/subprocess.*
shell_func = None
if isinstance(func, ast.Attribute):
if isinstance(func.value, ast.Name):
if func.value.id in self.os_aliases:
shell_func = f"os.{func.attr}"
elif func.value.id in self.subprocess_aliases:
shell_func = f"subprocess.{func.attr}"
elif isinstance(func, ast.Name):
# Check from-import aliases: from os import system; system(...)
shell_func = self.shell_exec_aliases.get(func.id)
if shell_func and shell_func in _SHELL_EXEC_FUNCS:
# Expand **kwargs dicts to inspect their keys
expanded_kwargs: dict[str, ast.AST] = {}
has_opaque_kwargs = False
for kw in node.keywords:
if kw.arg is not None:
expanded_kwargs[kw.arg] = kw.value
elif isinstance(kw.value, ast.Dict):
for k, v in zip(kw.value.keys, kw.value.values):
key = _extract_string_from_node(k) if k else None
if key is not None:
expanded_kwargs[key] = v
else:
has_opaque_kwargs = True
cmd_kw_values = [
v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS
]
all_call_args = list(node.args) + cmd_kw_values
blocked_in_args = _check_args_for_blocked(all_call_args)
if has_opaque_kwargs:
# Can't inspect dynamic **kwargs -- flag as unsafe
shell_escapes.append(
{
"type": "shell_escape_dynamic",
"line": node.lineno,
"description": (
f"{shell_func}() called with dynamic **kwargs"
),
}
)
elif blocked_in_args:
shell_escapes.append(
{
"type": "shell_escape",
"line": node.lineno,
"description": (
f"{shell_func}() invokes blocked command(s): "
f"{', '.join(sorted(blocked_in_args))}"
),
}
)
else:
# Only flag dynamic args for functions that interpret
# strings as shell commands, or when shell= might be
# enabled. Treat any non-literal-False shell= value
# as potentially True (conservative).
_STRING_SHELL_FUNCS = frozenset(
{
"os.system",
"os.popen",
"os.popen2",
"os.popen3",
"os.popen4",
"subprocess.getoutput",
"subprocess.getstatusoutput",
}
)
shell_node = expanded_kwargs.get("shell")
shell_safe = shell_node is None or (
isinstance(shell_node, ast.Constant)
and shell_node.value is False
)
# Dynamic shell-exec args (chr/format/concat bypasses).
if (
shell_func in _STRING_SHELL_FUNCS
or shell_func in _SHELL_EXEC_FUNCS
or not shell_safe
):
def _is_safe_literal(n):
if _extract_string_from_node(n) is not None:
return True
if isinstance(n, (ast.List, ast.Tuple)):
return all(
_extract_string_from_node(e) is not None
for e in n.elts
)
return False
has_non_literal = any(
not _is_safe_literal(a) for a in all_call_args
)
if has_non_literal:
shell_escapes.append(
{
"type": "shell_escape_dynamic",
"line": node.lineno,
"description": (
f"{shell_func}() called with non-literal "
f"shell command (potential shell escape)"
),
}
)
self.generic_visit(node)
def visit_ExceptHandler(self, node):
if self.loop_depth == 0:
self.generic_visit(node)
return
if node.type is None:
exception_catching.append(
{
"type": "bare_except_in_loop",
"line": node.lineno,
"description": "Bare except in loop catches TimeoutError and continues looping",
}
)
elif isinstance(node.type, ast.Name):
# Only flag BaseException and TimeoutError, NOT Exception.
# except Exception does not catch SystemExit or
# KeyboardInterrupt, so it cannot suppress timeout
# enforcement. Flagging Exception causes false positives
# on normal error-handling patterns.
if node.type.id in ("TimeoutError", "BaseException"):
exception_catching.append(
{
"type": f"catches_{node.type.id}_in_loop",
"line": node.lineno,
"description": f"Catches {node.type.id} in loop - may suppress timeout and continue",
}
)
elif isinstance(node.type, ast.Tuple):
for elt in node.type.elts:
if isinstance(elt, ast.Name):
if elt.id in ("TimeoutError", "BaseException"):
exception_catching.append(
{
"type": f"catches_{elt.id}_in_loop",
"line": node.lineno,
"description": f"Catches {elt.id} in loop - may suppress timeout and continue",
}
)
self.generic_visit(node)
visitor = SignalEscapeVisitor()
visitor.visit(tree)
if visitor.imports_signal and not signal_tampering:
warnings.append("Code imports 'signal' module - review manually for safety")
# Static host policy: block metadata hosts and any literal host outside
# the trusted allowlist; uploads blocked regardless of host. Dynamic hosts
# are caught by the bash blocklist instead.
network_calls: list[dict] = []
sensitive_file_reads: list[dict] = []
_NETWORK_FQ_PREFIXES = (
"socket.socket",
"socket.create_connection",
"socket.getaddrinfo",
"urllib.request.urlopen",
"urllib.request.urlretrieve",
"urllib3.",
"requests.get",
"requests.post",
"requests.put",
"requests.delete",
"requests.patch",
"requests.head",
"requests.request",
"requests.Session",
"http.client.HTTPConnection",
"http.client.HTTPSConnection",
"httpx.get",
"httpx.post",
"httpx.put",
"httpx.patch",
"httpx.delete",
"httpx.request",
"httpx.Client",
"httpx.AsyncClient",
"aiohttp.ClientSession",
)
_UPLOAD_HTTP_METHODS = (
"requests.post",
"requests.put",
"requests.patch",
"requests.delete",
"requests.request",
"httpx.post",
"httpx.put",
"httpx.patch",
"httpx.delete",
"httpx.request",
"urllib.request.urlopen",
"urllib.request.Request",
)
_UPLOAD_HF_FQ = (
"huggingface_hub.upload_file",
"huggingface_hub.upload_folder",
"huggingface_hub.upload_large_folder",
"huggingface_hub.create_commit",
)
_UPLOAD_HF_METHODS = frozenset(
{
"upload_file",
"upload_folder",
"upload_large_folder",
"create_commit",
}
)
# Cloud-metadata / link-local hosts.
_METADATA_HOST_LITERALS = {
"169.254.169.254",
"fd00:ec2::254",
"metadata.google.internal",
"metadata",
"metadata.tencentyun.com",
"100.100.100.200",
"100.100.100.110",
"169.254.170.2",
"169.254.170.23",
}
_METADATA_HOST_PREFIXES = (
"169.254.",
"100.64.",
)
# Allowlist kept explicit so each entry is auditable.
_TRUSTED_PUBLIC_HOST_LITERALS = frozenset(
{
# search
"www.google.com",
"google.com",
"www.bing.com",
"bing.com",
"duckduckgo.com",
"html.duckduckgo.com",
# encyclopedic / reference
"wikipedia.org",
"www.wikipedia.org",
"wikimedia.org",
"www.wikimedia.org",
"wikidata.org",
"www.wikidata.org",
"commons.wikimedia.org",
"www.britannica.com",
"openlibrary.org",
"www.openstreetmap.org",
# ML / dev / data
"huggingface.co",
"hf.co",
"github.com",
"api.github.com",
"raw.githubusercontent.com",
"gist.github.com",
"docs.github.com",
"pypi.org",
"files.pythonhosted.org",
"www.npmjs.com",
"registry.npmjs.org",
"crates.io",
"static.crates.io",
# docs
"docs.python.org",
"python.org",
"www.python.org",
"developer.mozilla.org",
"developer.apple.com",
"learn.microsoft.com",
"docs.docker.com",
"pytorch.org",
"docs.pytorch.org",
"tensorflow.org",
"www.tensorflow.org",
"numpy.org",
"pandas.pydata.org",
"scipy.org",
"scikit-learn.org",
"matplotlib.org",
"fastapi.tiangolo.com",
"starlette.io",
# academic
"arxiv.org",
"export.arxiv.org",
"scholar.google.com",
"openreview.net",
"semanticscholar.org",
"www.semanticscholar.org",
"biorxiv.org",
"www.biorxiv.org",
"medrxiv.org",
"www.medrxiv.org",
"pubmed.ncbi.nlm.nih.gov",
"www.ncbi.nlm.nih.gov",
# Q&A / community
"stackoverflow.com",
"stackexchange.com",
"askubuntu.com",
"superuser.com",
"serverfault.com",
# standards
"www.w3.org",
"tools.ietf.org",
"datatracker.ietf.org",
"www.rfc-editor.org",
# reputable news
"www.bbc.com",
"www.bbc.co.uk",
"www.reuters.com",
"apnews.com",
"www.nature.com",
"www.science.org",
# government / open data
"data.gov",
"catalog.data.gov",
"www.census.gov",
"www.nasa.gov",
"data.nasa.gov",
"www.cdc.gov",
"www.nih.gov",
"www.who.int",
# weather / time
"api.weather.gov",
"worldtimeapi.org",
}
)
_TRUSTED_PUBLIC_HOST_SUFFIXES = (
".wikipedia.org",
".wikimedia.org",
".wiktionary.org",
".wikibooks.org",
".wikiquote.org",
".wikisource.org",
".wikiversity.org",
".wikivoyage.org",
".stackexchange.com",
".hf.co",
".huggingface.co",
".githubusercontent.com",
".github.io",
".arxiv.org",
".readthedocs.io",
".readthedocs.org",
)
_SENSITIVE_FILE_PREFIXES = (
"/etc/passwd",
"/etc/shadow",
"/etc/sudoers",
"/etc/ssh/",
)
_SENSITIVE_FILE_RE = re.compile(
r"^/proc/(?:self|\d+)/(?:environ|cmdline|task/\d+/environ)$"
)
def _normalize_host(host: str) -> str:
if not host:
return ""
h = host.strip().lower().rstrip(".")
if "@" in h:
h = h.split("@", 1)[1]
if h.startswith("[") and "]" in h:
h = h[1 : h.index("]")]
elif h.count(":") == 1:
h = h.split(":", 1)[0]
return h
def _is_metadata_host(host: str) -> bool:
h = _normalize_host(host)
if not h:
return False
if h in _METADATA_HOST_LITERALS:
return True
if any(h.startswith(p) for p in _METADATA_HOST_PREFIXES):
return True
return False
def _is_trusted_host(host: str) -> bool:
h = _normalize_host(host)
if not h:
return False
if h in _TRUSTED_PUBLIC_HOST_LITERALS:
return True
return any(h.endswith(s) for s in _TRUSTED_PUBLIC_HOST_SUFFIXES)
def _call_is_upload_shape(node: ast.Call, fq: str) -> bool:
"""True for statically obvious upload shapes (files=, data=open(), bytes literal)."""
if fq in _UPLOAD_HF_FQ:
return True
if fq not in _UPLOAD_HTTP_METHODS:
return False
for kw in node.keywords or []:
if kw.arg == "files":
return True
if kw.arg == "data":
v = kw.value
if (
isinstance(v, ast.Call)
and isinstance(v.func, ast.Name)
and v.func.id == "open"
):
return True
if isinstance(v, ast.Constant) and isinstance(
v.value, (bytes, bytearray)
):
return True
return False
# Bare method-name fallback (`x.upload_file(...)`) is intentionally fuzzy,
# but should only fire when huggingface_hub / hf_api is actually imported
# somewhere in the snippet -- otherwise paramiko.upload_file, boto3
# create_commit, etc. hit a false positive. We pre-scan for the imports.
_HF_IMPORT_MODULES = (
"huggingface_hub",
"hf_api",
"huggingface_hub.hf_api",
)
def _module_has_hf_import(tree: ast.AST) -> bool:
for n in ast.walk(tree):
if isinstance(n, ast.Import):
for alias in n.names:
if alias.name.split(".", 1)[0] in _HF_IMPORT_MODULES:
return True
elif isinstance(n, ast.ImportFrom):
root = (n.module or "").split(".", 1)[0]
if root in _HF_IMPORT_MODULES:
return True
elif isinstance(n, ast.Call) and n.args:
# __import__('huggingface_hub'), importlib.import_module('huggingface_hub'),
# and bare import_module('huggingface_hub') (via `from importlib import ...`).
arg0 = n.args[0]
if not (isinstance(arg0, ast.Constant) and isinstance(arg0.value, str)):
continue
if arg0.value.split(".", 1)[0] not in _HF_IMPORT_MODULES:
continue
func = n.func
if isinstance(func, ast.Name) and func.id in {
"__import__",
"import_module",
}:
return True
if isinstance(func, ast.Attribute) and func.attr == "import_module":
return True
return False
_hf_in_scope = _module_has_hf_import(tree)
def _method_call_hf_upload_name(node: ast.Call) -> str | None:
"""Return the HF upload method name (`upload_file`, ...) or None.
Catches `HfApi().upload_file(...)` (Attribute) and
`from huggingface_hub import upload_file; upload_file(...)` (Name).
The bare-name branch fires only when an HF import is in scope, mirroring
the Attribute branch's gating so paramiko/boto3 do not false-positive.
"""
if not _hf_in_scope:
return None
f = node.func
if isinstance(f, ast.Attribute) and f.attr in _UPLOAD_HF_METHODS:
return f.attr
if isinstance(f, ast.Name) and f.id in _UPLOAD_HF_METHODS:
return f.id
return None
# Kwargs that ship a credential over the wire. Sandbox env strips HF_TOKEN
# / WANDB_API_KEY / AWS_* up front, so any value here is hard-coded or
# lifted from the parent process.
_HF_SENSITIVE_KWARGS = frozenset(
{
"token",
"hf_token",
"api_token",
"api_key",
"auth_token",
"access_token",
"password",
"secret",
}
)
def _is_os_environ(node: ast.AST) -> bool:
return (
isinstance(node, ast.Attribute)
and node.attr == "environ"
and isinstance(node.value, ast.Name)
and node.value.id == "os"
)
def _reads_env_or_secret(node: ast.AST | None) -> bool:
"""True if any node in the subtree resolves to an env / process read.
Walking the subtree (not just the root) means wrapper calls like
`str(os.environ)`, `json.dumps(os.environ)`, or
`'-'.join(os.environ.values())` are caught too.
Covers: `os.environ`, `os.environ[K]`, `os.environ.get(K)`, `os.getenv(K)`,
bare `getenv(K)` (after `from os import getenv`), and
`subprocess.{run,check_output,Popen,getoutput,getstatusoutput}` which
the LLM could use to lift parent env via `printenv` / `env` / `set`.
"""
if node is None:
return False
for sub in ast.walk(node):
if _is_os_environ(sub):
return True
if isinstance(sub, ast.Call):
f = sub.func
if isinstance(f, ast.Attribute):
if (
f.attr in {"getenv", "getenvb"}
and isinstance(f.value, ast.Name)
and f.value.id == "os"
):
return True
if (
f.attr
in {
"check_output",
"run",
"Popen",
"getoutput",
"getstatusoutput",
}
and isinstance(f.value, ast.Name)
and f.value.id in {"subprocess", "commands"}
):
return True
if isinstance(f, ast.Name) and f.id in {"getenv", "getenvb"}:
return True
return False
def _is_safe_relative_path(path: str) -> bool:
"""Relative path with no leading `/`, `~`, drive letter, or `..` segments."""
if not isinstance(path, str) or not path:
return False
if path[0] in ("/", "\\", "~"):
return False
if len(path) >= 2 and path[1] == ":":
return False
return ".." not in path.replace("\\", "/").split("/")
def _path_arg_is_sandbox_local(node: ast.AST | None) -> bool:
"""Whether the path argument resolves to a sandbox-local literal."""
if node is None:
return False
if isinstance(node, ast.Constant) and isinstance(
node.value, (bytes, bytearray)
):
return True # inline bytes, no file access
if isinstance(node, ast.Constant) and isinstance(node.value, str):
return _is_safe_relative_path(node.value)
if isinstance(node, ast.Call):
f = node.func
is_open = (isinstance(f, ast.Name) and f.id == "open") or (
isinstance(f, ast.Attribute) and f.attr == "open"
)
if is_open and node.args:
a0 = node.args[0]
return (
isinstance(a0, ast.Constant)
and isinstance(a0.value, str)
and _is_safe_relative_path(a0.value)
)
return False
def _hf_upload_violation(node: ast.Call, method_name: str) -> str | None:
"""Inspect an HF upload call; return a violation reason or None.
Policy: HF uploads are allowed only when (a) no sensitive kwarg is set,
(b) no positional / keyword value reads `os.environ` or related env
readers, and (c) the path argument is a sandbox-local literal -- a
relative string with no `..`, an `open(<literal>)`, or inline bytes.
Dynamic / variable paths are rejected; the policy cannot prove safety
statically and the cost of a wrong-allow is a credential exfiltration.
"""
for kw in node.keywords or []:
if kw.arg in _HF_SENSITIVE_KWARGS:
return (
f"HF upload {kw.arg}= cannot be set from sandboxed code; "
"uploads run with the sandbox identity only"
)
all_values = list(node.args or []) + [kw.value for kw in (node.keywords or [])]
for v in all_values:
if _reads_env_or_secret(v):
return (
"HF upload cannot include os.environ / os.getenv / subprocess "
"env reads; secrets and tokens must not be exfiltrated"
)
if method_name == "create_commit":
for kw in node.keywords or []:
if kw.arg == "operations" and isinstance(kw.value, ast.List):
for elt in kw.value.elts:
if isinstance(elt, ast.Call):
inner = _hf_upload_violation(elt, "upload_file")
if inner:
return inner
return None
path_node: ast.AST | None = node.args[0] if node.args else None
for kw in node.keywords or []:
if kw.arg in ("path_or_fileobj", "folder_path"):
path_node = kw.value
break
if not _path_arg_is_sandbox_local(path_node):
return (
"HF upload path must be a sandbox-local relative-path literal "
"(no absolute paths, no '..' segments, no dynamic expressions)"
)
return None
class NetworkAndIoVisitor(ast.NodeVisitor):
def visit_Call(self, node):
parts: list[str] = []
cur = node.func
while isinstance(cur, ast.Attribute):
parts.insert(0, cur.attr)
cur = cur.value
if isinstance(cur, ast.Name):
parts.insert(0, cur.id)
fq = ".".join(parts) if parts else ""
hf_upload_name = _method_call_hf_upload_name(node)
if hf_upload_name is not None:
violation = _hf_upload_violation(node, hf_upload_name)
if violation is not None:
network_calls.append(
{
"type": "upload_blocked",
"line": getattr(node, "lineno", -1),
"description": f"Blocked: {violation}",
}
)
# Direct sock.connect((host, port)) bypasses the FQ-prefix branch below.
if (
isinstance(node.func, ast.Attribute)
and node.func.attr == "connect"
and node.args
):
a0 = node.args[0]
host_lit = None
if isinstance(a0, ast.Tuple) and a0.elts:
e0 = a0.elts[0]
if isinstance(e0, ast.Constant) and isinstance(e0.value, str):
host_lit = e0.value
elif isinstance(a0, ast.Constant) and isinstance(a0.value, str):
host_lit = a0.value
if host_lit:
if _is_metadata_host(host_lit):
network_calls.append(
{
"type": "metadata_host_blocked",
"line": getattr(node, "lineno", -1),
"description": "Blocked: cloud-metadata host",
}
)
elif not _is_trusted_host(host_lit):
network_calls.append(
{
"type": "untrusted_host_blocked",
"line": getattr(node, "lineno", -1),
"description": (
"Blocked: host not in sandbox allowlist; "
"use an allowed informational source"
),
}
)
if fq and any(fq.startswith(p) for p in _NETWORK_FQ_PREFIXES):
# 1) Upload-shape check (host-independent).
if _call_is_upload_shape(node, fq):
network_calls.append(
{
"type": "upload_blocked",
"line": getattr(node, "lineno", -1),
"description": (
"Blocked: file upload disallowed in sandbox"
),
}
)
# 2) Extract literal host (URL string or (host, port) tuple).
host_arg = None
url_arg = None
if node.args:
a0 = node.args[0]
if isinstance(a0, ast.Constant) and isinstance(a0.value, str):
url_arg = a0.value
elif isinstance(a0, ast.Tuple) and a0.elts:
e0 = a0.elts[0]
if isinstance(e0, ast.Constant) and isinstance(e0.value, str):
host_arg = e0.value
if url_arg and host_arg is None:
m = re.match(r"^\w+://([^/?#]+)", url_arg)
if m:
host_arg = m.group(1)
if host_arg:
if _is_metadata_host(host_arg):
network_calls.append(
{
"type": "metadata_host_blocked",
"line": getattr(node, "lineno", -1),
"description": "Blocked: cloud-metadata host",
}
)
elif not _is_trusted_host(host_arg):
network_calls.append(
{
"type": "untrusted_host_blocked",
"line": getattr(node, "lineno", -1),
"description": (
"Blocked: host not in sandbox allowlist; "
"use an allowed informational source"
),
}
)
is_open_call = (
(isinstance(node.func, ast.Name) and node.func.id == "open")
or fq in ("io.open", "pathlib.Path.open")
or fq.endswith(".open")
)
if is_open_call and node.args:
a0 = node.args[0]
path_lit = None
if isinstance(a0, ast.Constant) and isinstance(a0.value, str):
path_lit = a0.value
if path_lit:
flagged = False
if any(path_lit.startswith(p) for p in _SENSITIVE_FILE_PREFIXES):
flagged = True
elif _SENSITIVE_FILE_RE.match(path_lit):
flagged = True
if flagged:
sensitive_file_reads.append(
{
"type": "sensitive_file_read",
"line": getattr(node, "lineno", -1),
"description": (
f"open({path_lit!r}) targets a host identity / "
"credential file; sandboxed code may not read it"
),
}
)
self.generic_visit(node)
NetworkAndIoVisitor().visit(tree)
is_safe = (
len(signal_tampering) == 0
and len(exception_catching) == 0
and len(shell_escapes) == 0
and len(network_calls) == 0
and len(sensitive_file_reads) == 0
)
return is_safe, {
"signal_tampering": signal_tampering,
"exception_catching": exception_catching,
"shell_escapes": shell_escapes,
"network_calls": network_calls,
"sensitive_file_reads": sensitive_file_reads,
"warnings": warnings,
}
def _check_code_safety(code: str) -> str | None:
"""Validate code safety via static analysis.
Returns an error message string if the code is unsafe, or None if OK.
"""
safe, info = _check_signal_escape_patterns(code)
if not safe:
# SyntaxError from ast.parse -- let these through so the subprocess
# produces a normal Python traceback instead of a misleading
# "unsafe code detected" message.
if info.get("error"):
return None
reasons = [
item.get("description", "") for item in info.get("signal_tampering", [])
]
shell_reasons = [
item.get("description", "") for item in info.get("shell_escapes", [])
]
exception_reasons = [
item.get("description", "") for item in info.get("exception_catching", [])
]
network_reasons = [
item.get("description", "") for item in info.get("network_calls", [])
]
file_reasons = [
item.get("description", "") for item in info.get("sensitive_file_reads", [])
]
all_reasons = [
r
for r in reasons
+ shell_reasons
+ exception_reasons
+ network_reasons
+ file_reasons
if r
]
if all_reasons:
return (
f"Error: unsafe code detected ({'; '.join(all_reasons)}). "
f"Please remove unsafe patterns from your code."
)
return None
def _kill_process_tree(proc) -> None:
"""SIGKILL the setsid process group; fall back to single-pid kill."""
if proc.poll() is not None:
return
try:
pgid = os.getpgid(proc.pid)
except (ProcessLookupError, PermissionError):
pgid = None
if pgid is not None:
try:
os.killpg(pgid, signal.SIGKILL)
return
except (ProcessLookupError, PermissionError):
pass
try:
proc.kill()
except (ProcessLookupError, PermissionError):
pass
def _cancel_watcher(proc, cancel_event, poll_interval = 0.2):
"""Daemon thread that kills a process when cancel_event is set."""
while proc.poll() is None:
if cancel_event is not None and cancel_event.is_set():
_kill_process_tree(proc)
return
cancel_event.wait(poll_interval) if cancel_event else None
def _truncate(text: str, limit: int = _MAX_OUTPUT_CHARS) -> str:
if len(text) > limit:
return text[:limit] + f"\n\n... (truncated, {len(text)} chars total)"
return text
def _python_exec(
code: str,
cancel_event = None,
timeout: int = _EXEC_TIMEOUT,
session_id: str | None = None,
) -> str:
"""Execute Python code in a subprocess sandbox."""
if not code or not code.strip():
return "No code provided."
# Validate imports and code safety
error = _check_code_safety(code)
if error:
return error
tmp_path = None
workdir = _get_workdir(session_id)
# Snapshot image mtimes so we detect both new and overwritten files.
_before: dict[str, int] = {}
if os.path.isdir(workdir):
for _name in os.listdir(workdir):
if os.path.splitext(_name)[1].lower() in _IMAGE_EXTS:
_p = os.path.join(workdir, _name)
if os.path.isfile(_p):
try:
_before[_name] = os.stat(_p).st_mtime_ns
except OSError:
pass
try:
fd, tmp_path = tempfile.mkstemp(
suffix = ".py", prefix = "studio_exec_", dir = workdir
)
with os.fdopen(fd, "w") as f:
f.write(code)
safe_env = _build_safe_env(workdir)
popen_kwargs = dict(
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
cwd = workdir,
env = safe_env,
)
if sys.platform != "win32":
popen_kwargs["preexec_fn"] = _sandbox_preexec
else:
popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
proc = subprocess.Popen([sys.executable, tmp_path], **popen_kwargs)
# Spawn cancel watcher if we have a cancel event
if cancel_event is not None:
watcher = threading.Thread(
target = _cancel_watcher, args = (proc, cancel_event), daemon = True
)
watcher.start()
try:
output, _ = proc.communicate(timeout = timeout)
except subprocess.TimeoutExpired:
_kill_process_tree(proc)
try:
proc.communicate(timeout = 5)
except subprocess.TimeoutExpired:
pass
return _truncate(f"Execution timed out after {timeout} seconds.")
if cancel_event is not None and cancel_event.is_set():
return "Execution cancelled."
result = output or ""
if proc.returncode != 0:
result = f"Exit code {proc.returncode}:\n{result}"
result = _truncate(result) if result.strip() else "(no output)"
# Detect new or overwritten image files and append sentinel for frontend
if session_id and os.path.isdir(workdir):
new_images = []
for _name in os.listdir(workdir):
if os.path.splitext(_name)[1].lower() not in _IMAGE_EXTS:
continue
_p = os.path.join(workdir, _name)
if not os.path.isfile(_p):
continue
try:
_mtime = os.stat(_p).st_mtime_ns
except OSError:
continue
if _name not in _before or _mtime != _before[_name]:
new_images.append(_name)
if new_images:
import json as _json
result += f"\n__IMAGES__:{_json.dumps(sorted(new_images))}"
return result
except Exception as e:
return f"Execution error: {e}"
finally:
if tmp_path and os.path.exists(tmp_path):
try:
os.unlink(tmp_path)
except OSError:
pass
def _bash_exec(
command: str,
cancel_event = None,
timeout: int = _EXEC_TIMEOUT,
session_id: str | None = None,
) -> str:
"""Execute a bash command in a subprocess sandbox."""
if not command or not command.strip():
return "No command provided."
# Block dangerous commands (shlex + regex based)
blocked = _find_blocked_commands(command)
if blocked:
return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
try:
workdir = _get_workdir(session_id)
safe_env = _build_safe_env(workdir)
popen_kwargs = dict(
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
cwd = workdir,
env = safe_env,
)
if sys.platform != "win32":
popen_kwargs["preexec_fn"] = _sandbox_preexec
else:
popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
proc = subprocess.Popen(_get_shell_cmd(command), **popen_kwargs)
if cancel_event is not None:
watcher = threading.Thread(
target = _cancel_watcher, args = (proc, cancel_event), daemon = True
)
watcher.start()
try:
output, _ = proc.communicate(timeout = timeout)
except subprocess.TimeoutExpired:
_kill_process_tree(proc)
try:
proc.communicate(timeout = 5)
except subprocess.TimeoutExpired:
pass
return _truncate(f"Execution timed out after {timeout} seconds.")
if cancel_event is not None and cancel_event.is_set():
return "Execution cancelled."
result = output or ""
if proc.returncode != 0:
result = f"Exit code {proc.returncode}:\n{result}"
return _truncate(result) if result.strip() else "(no output)"
except Exception as e:
return f"Execution error: {e}"