File size: 2,762 Bytes
c641d5f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """Constrained execution for small benchmark Python attachments."""
from __future__ import annotations
import ast
import subprocess
import sys
from pathlib import Path
_ALLOWED_IMPORTS = {
"cmath",
"collections",
"decimal",
"fractions",
"functools",
"itertools",
"math",
"operator",
"random",
"re",
"statistics",
"time",
}
_BLOCKED_CALLS = {
"breakpoint",
"compile",
"eval",
"exec",
"input",
"open",
"__import__",
}
def _validate(source: str) -> None:
tree = ast.parse(source)
for node in ast.walk(tree):
if isinstance(node, (ast.Import, ast.ImportFrom)):
names = (
[a.name.split(".")[0] for a in node.names]
if isinstance(node, ast.Import)
else [(node.module or "").split(".")[0]]
)
if any(name not in _ALLOWED_IMPORTS for name in names):
raise ValueError(
f"Blocked import in Python attachment: {', '.join(names)}"
)
if (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id in _BLOCKED_CALLS
):
raise ValueError(f"Blocked call in Python attachment: {node.func.id}")
if isinstance(node, ast.Attribute) and node.attr.startswith("__"):
raise ValueError("Dunder attribute access is blocked in Python attachments")
def execute_python_file(path: str | Path, timeout: int = 15) -> str:
"""Run a validated Python file in isolated interpreter mode and return stdout."""
file_path = Path(path).resolve()
if file_path.suffix.lower() != ".py" or not file_path.is_file():
raise ValueError("Expected an existing .py file")
source = file_path.read_text(encoding="utf-8")
_validate(source)
completed = subprocess.run(
[
sys.executable,
"-I",
"-c",
(
"import runpy,sys,time; "
"time.sleep=lambda _seconds: None; "
"runpy.run_path(sys.argv[1],run_name='__main__')"
),
str(file_path),
],
cwd=file_path.parent,
capture_output=True,
text=True,
timeout=timeout,
check=False,
env={"PYTHONIOENCODING": "utf-8"},
)
if completed.returncode != 0:
error = completed.stderr.strip()[-2000:]
raise RuntimeError(f"Python attachment exited {completed.returncode}: {error}")
output = completed.stdout.strip()
if not output:
raise RuntimeError("Python attachment produced no stdout")
return output
|