execute.py

"""Sandboxed execution + validation of generated plotting code."""

import multiprocessing as mp
import re
import matplotlib

matplotlib.use("Agg")  # Headless mode
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

try:
    import seaborn as sns
except ImportError:
    sns = None

from schemas import GENERATORS

# ---------------------------------------------------------------------------
# Execution guard
# ---------------------------------------------------------------------------
# The code we exec() here is machine-generated (Gemini teacher, and later the
# fine-tuned student). It is NOT trusted. Running it with the full builtins and
# unrestricted imports would let a hallucinated snippet run `import os;
# os.system(...)`, delete files, or open network sockets on this machine.
#
# This is a *proportionate* guard against accidental / hallucinated harm, not a
# hardened sandbox: it removes the obvious escape hatches (open/eval/exec/compile
# and imports of os/subprocess/socket/etc.) while still allowing normal plotting
# code. For genuinely adversarial code, run this inside a container or VM.

# Top-level modules a plotting snippet is allowed to import.
_ALLOWED_IMPORTS = {
    "matplotlib", "mpl_toolkits", "pandas", "numpy", "seaborn",
    "scipy", "math", "statistics", "datetime", "collections",
    "itertools", "functools", "warnings", "random", "cycler",
    "numbers", "decimal", "fractions",
}

# Builtins safe to expose. Deliberately omits open, eval, exec, compile, input,
# __import__ (replaced below), globals, locals, vars, exit, quit, help,
# breakpoint, memoryview.
_SAFE_BUILTIN_NAMES = (
    "abs", "all", "any", "bool", "bytes", "callable", "chr", "classmethod",
    "complex", "dict", "divmod", "enumerate", "filter", "float", "format",
    "frozenset", "getattr", "hasattr", "hash", "hex", "int", "isinstance",
    "issubclass", "iter", "len", "list", "map", "max", "min", "next", "object",
    "oct", "ord", "pow", "print", "property", "range", "repr", "reversed",
    "round", "set", "setattr", "slice", "sorted", "staticmethod", "str", "sum",
    "super", "tuple", "type", "zip", "True", "False", "None",
    "__build_class__",  # needed for class definitions in generated code
)


def _guarded_import(name, globals=None, locals=None, fromlist=(), level=0):
    """Allow imports only for the whitelisted data/plotting modules."""
    root = name.split(".")[0]
    if root not in _ALLOWED_IMPORTS:
        raise ImportError(f"import of '{name}' is blocked in the execution sandbox")
    return __import__(name, globals, locals, fromlist, level)


def _safe_builtins() -> dict:
    import builtins as _b
    safe = {n: getattr(_b, n) for n in _SAFE_BUILTIN_NAMES if hasattr(_b, n)}
    safe["__import__"] = _guarded_import
    return safe


GENERATOR_MAP = (
    {g.__name__: g if callable(g) else g for g in GENERATORS}
    if isinstance(GENERATORS, (list, tuple, set))
    else GENERATORS
)


def clean_code_string(code: str) -> str:
    """Strip markdown code blocks from generated code safely."""
    if not code:
        return ""

    code = code.strip()
    bt3 = chr(96) * 3  # Dynamically generates triple backticks to avoid UI glitches

    # Extract content inside markdown python code block if present
    pattern = bt3 + r"(?:python)?\s*\n?(.*?)\n?" + bt3
    match = re.search(pattern, code, flags=re.DOTALL | re.IGNORECASE)
    
    if match:
        cleaned = match.group(1).strip()
    else:
        # Fallback: strip leading/trailing backtick lines
        lines = code.splitlines()
        if lines and lines[0].strip().startswith(bt3):
            lines = lines[1:]
        if lines and lines[-1].strip().startswith(bt3):
            lines = lines[:-1]
        cleaned = "\n".join(lines).strip()

    return cleaned


def _run(gen_name: str, code: str, q: mp.Queue, seed: int = 0) -> None:
    # Generated code triggers a lot of library deprecation chatter (seaborn
    # palette/hue, pandas futures). It is not our code and we cannot fix it, so
    # keep it out of the validation log -- real failures come back via the queue.
    import warnings
    warnings.filterwarnings("ignore")
    try:
        if callable(gen_name):
            gen_func = gen_name
        elif gen_name in GENERATOR_MAP:
            gen_func = GENERATOR_MAP[gen_name]
        else:
            q.put(("error", f"Generator '{gen_name}' not found."))
            return

        clean_code = clean_code_string(code)
        if not clean_code:
            q.put(("error", "Empty code block after stripping formatting."))
            return

        # Same (generator, seed) the teacher saw -> the code is validated against
        # the exact DataFrame whose preview is stored on the record.
        df = gen_func(seed=seed)

        ns = {
            "df": df,
            "plt": plt,
            "pd": pd,
            "np": np,
            "matplotlib": matplotlib,
            "__builtins__": _safe_builtins(),  # restricted: no os/open/eval/etc.
        }
        if sns is not None:
            ns["sns"] = sns

        plt.close("all")
        exec(clean_code, ns)

        fignums = plt.get_fignums()
        if not fignums:
            q.put(("no_figure", "No matplotlib figure was created."))
            return

        fig = plt.gcf()
        has_axes = len(fig.axes) > 0 and any(len(ax.get_children()) > 0 for ax in fig.axes)
        plt.close("all")

        if has_axes:
            q.put(("ok", None))
        else:
            q.put(("no_figure", "Figure created but contains no axes or visual elements."))
    except Exception as e:
        plt.close("all")
        q.put(("error", f"{type(e).__name__}: {str(e)}"))


def validate(gen_name: str, code: str, timeout: int = 10, seed: int = 0) -> tuple[bool, str | None]:
    if not code:
        return False, "empty code"

    q = mp.Queue()
    p = mp.Process(target=_run, args=(gen_name, code, q, seed))
    p.start()
    p.join(timeout)

    if p.is_alive():
        p.terminate()
        p.join(timeout=2)
        if p.is_alive():
            p.kill()
            p.join()
        return False, f"Timeout after {timeout} seconds"

    if q.empty():
        return False, "Process terminated unexpectedly with no output"

    status, detail = q.get()
    return status == "ok", detail