Agent Security Basics

python -m venv .venv && source .venv/bin/activate
pip install pydantic

THREATS = [
    "read host secrets via open('/etc/passwd') or Path.home()",
    "exfiltrate data with urllib or socket",
    "fork-bomb or infinite loop",
    "import os and call system()",
]
print("sandbox must block:", *THREATS, sep="\n- ")

from pydantic import BaseModel, Field
 
 
class RunPythonArgs(BaseModel):
    source: str = Field(
        description="Short pure Python to compute a result. No I/O, no imports."
    )
 
 
tools = [{
    "type": "function",
    "function": {
        "name": "run_python_sandboxed",
        "description": "Run a short pure Python snippet and return stdout/result. No network or files.",
        "parameters": RunPythonArgs.model_json_schema(),
    },
}]

import ast
from typing import Any
 
 
ALLOWED_BUILTINS = {
    "abs": abs,
    "min": min,
    "max": max,
    "sum": sum,
    "len": len,
    "range": range,
    "enumerate": enumerate,
    "sorted": sorted,
    "round": round,
    "float": float,
    "int": int,
    "str": str,
    "bool": bool,
    "list": list,
    "dict": dict,
    "tuple": tuple,
    "print": print,
}
 
 
def run_sandboxed(source: str, *, timeout_s: float = 1.0) -> dict[str, Any]:
    # Timeout needs threads/processes in real systems; kept simple here.
    tree = ast.parse(source, mode="exec")
    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            raise PermissionError("imports are not allowed")
        if isinstance(node, ast.Attribute) and isinstance(node.attr, str):
            if node.attr.startswith("_"):
                raise PermissionError("private attributes are not allowed")
 
    stdout: list[str] = []
 
    def _print(*args, **kwargs):
        stdout.append(" ".join(str(a) for a in args))
 
    glb = {"__builtins__": {**ALLOWED_BUILTINS, "print": _print}}
    loc: dict[str, Any] = {}
    exec(compile(tree, "<agent>", "exec"), glb, loc)  # noqa: S102 - intentional sandbox demo
    result = loc.get("result", stdout[-1] if stdout else None)
    return {"ok": True, "result": result, "stdout": stdout}

import json
 
 
def handle_run_python_sandboxed(arguments_json: str) -> str:
    try:
        args = RunPythonArgs.model_validate_json(arguments_json)
        if len(args.source) > 4_000:
            raise ValueError("source too long")
        out = run_sandboxed(args.source)
        return json.dumps(out)
    except Exception as e:
        return json.dumps({"ok": False, "error": type(e).__name__, "message": str(e)[:200]})
 
 
print(handle_run_python_sandboxed('{"source": "result = sum(range(10))"}'))
print(handle_run_python_sandboxed('{"source": "import os"}'))

cases = [
    "result = open('secret.txt').read()",
    "import socket",
    "result = __import__('os').system('id')",
]
 
for src in cases:
    obs = handle_run_python_sandboxed(json.dumps({"source": src}))
    assert json.loads(obs)["ok"] is False, src
print("negative tests passed")

import multiprocessing as mp
from typing import Any
 
 
def _worker(source: str, conn):
    try:
        conn.send(run_sandboxed(source))
    except Exception as e:
        conn.send({"ok": False, "error": type(e).__name__, "message": str(e)[:200]})
    finally:
        conn.close()
 
 
def run_with_timeout(source: str, timeout_s: float = 1.0) -> dict[str, Any]:
    parent, child = mp.Pipe(duplex=False)
    proc = mp.Process(target=_worker, args=(source, child), daemon=True)
    proc.start()
    child.close()
    proc.join(timeout_s)
    if proc.is_alive():
        proc.kill()
        proc.join()
        return {"ok": False, "error": "Timeout", "message": f"exceeded {timeout_s}s"}
    if parent.poll():
        return parent.recv()
    return {"ok": False, "error": "Empty", "message": "no result"}
 
 
print(run_with_timeout("result = 1 + 1"))
print(run_with_timeout("while True: pass", timeout_s=0.3))

import os
from contextlib import contextmanager
 
 
@contextmanager
def scrubbed_env(allow: set[str] | None = None):
    allow = allow or set()
    saved = dict(os.environ)
    try:
        os.environ.clear()
        for k, v in saved.items():
            if k in allow:
                os.environ[k] = v
        yield
    finally:
        os.environ.clear()
        os.environ.update(saved)
 
 
# Agent code path should not see OPENAI_API_KEY / DB URLs
with scrubbed_env(allow={"PATH", "LANG"}):
    assert "OPENAI_API_KEY" not in os.environ
print("credentials not visible to tool env")

from typing import Callable
 
 
REGISTRY: dict[str, Callable[[str], str]] = {
    "run_python_sandboxed": handle_run_python_sandboxed,
}
 
 
def run_agent_turn(tool_name: str, arguments_json: str, *, max_turns: int = 3) -> list[dict]:
    """Demo host: only allowlisted tools, hard turn cap."""
    if tool_name not in REGISTRY:
        return [{"ok": False, "error": "ToolNotAllowed"}]
    observations = []
    for turn in range(max_turns):
        observations.append({
            "turn": turn,
            "tool": tool_name,
            "observation": json.loads(REGISTRY[tool_name](arguments_json)),
        })
        # Real agents re-call the model here; we stop after one successful compute.
        if observations[-1]["observation"].get("ok"):
            break
    return observations
 
 
print(run_agent_turn("run_python_sandboxed", '{"source": "result = 6 * 7"}'))
print(run_agent_turn("run_shell", '{"cmd": "ls"}'))  # denied