app.py

import io
import sys
import time
import ast
import traceback
from multiprocessing import Process, Queue

try:
    import resource  # Not available on Windows, but Spaces run Linux.
except Exception:  # pragma: no cover
    resource = None

import gradio as gr

# -----------------------------
# Sandbox utilities
# -----------------------------
SAFE_BUILTINS = {
    "abs": abs,
    "all": all,
    "any": any,
    "bool": bool,
    "bytes": bytes,
    "callable": callable,
    "chr": chr,
    "complex": complex,
    "dict": dict,
    "dir": dir,
    "divmod": divmod,
    "enumerate": enumerate,
    "filter": filter,
    "float": float,
    "format": format,
    "hash": hash,
    "help": help,
    "hex": hex,
    "int": int,
    "isinstance": isinstance,
    "issubclass": issubclass,
    "iter": iter,
    "len": len,
    "list": list,
    "map": map,
    "max": max,
    "min": min,
    "next": next,
    "object": object,
    "oct": oct,
    "ord": ord,
    "pow": pow,
    "print": print,
    "range": range,
    "repr": repr,
    "reversed": reversed,
    "round": round,
    "set": set,
    "slice": slice,
    "sorted": sorted,
    "str": str,
    "sum": sum,
    "tuple": tuple,
    "type": type,
    "zip": zip,
}

DANGEROUS_NAMES = {
    "open", "compile", "eval", "exec", "__import__", "input",
    "globals", "locals", "vars", "exit", "quit",
}

DANGEROUS_MODULE_TOKENS = {
    "os", "sys", "subprocess", "socket", "shutil", "ctypes", "signal", "resource",
}


def _lint_code_for_danger(code: str, allow_imports: bool) -> None:
    """Lightweight static checks; raises ValueError on disallowed constructs."""
    try:
        tree = ast.parse(code, mode="exec")
    except SyntaxError as e:
        raise ValueError(f"SyntaxError: {e}")

    for node in ast.walk(tree):
        # Disallow Attribute or Name usage of dunder or dangerous names
        if isinstance(node, ast.Name):
            if node.id in DANGEROUS_NAMES:
                raise ValueError(f"Use of `{node.id}` is not allowed in this sandbox.")
            if not allow_imports and node.id in DANGEROUS_MODULE_TOKENS:
                raise ValueError(
                    f"Importing/using `{node.id}` is not allowed in this sandbox. Enable 'Allow imports' to proceed (unsafe)."
                )
        if isinstance(node, ast.Attribute):
            if node.attr.startswith("__"):
                raise ValueError("Access to dunder attributes is blocked in this sandbox.")
        if isinstance(node, ast.Import) or isinstance(node, ast.ImportFrom):
            if not allow_imports:
                raise ValueError("Imports are disabled. Tick 'Allow imports (unsafe)' to enable.")


def _restricted_import(*args, **kwargs):
    raise ImportError("Importing modules is disabled in this sandbox.")


def _apply_limits(mem_mb: int, cpu_seconds: int):
    if resource is None:
        return
    # Limit address space (virtual memory)
    if mem_mb and mem_mb > 0:
        bytes_limit = int(mem_mb) * 1024 * 1024
        resource.setrlimit(resource.RLIMIT_AS, (bytes_limit, bytes_limit))
    # Limit CPU seconds
    if cpu_seconds and cpu_seconds > 0:
        resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds))


def _worker(code: str, stdin_data: str, allow_imports: bool, mem_limit_mb: int, cpu_limit_s: int, q: Queue):
    stdout_buf, stderr_buf = io.StringIO(), io.StringIO()
    start = time.time()
    try:
        _apply_limits(mem_limit_mb, cpu_limit_s)

        # Prepare I/O redirection
        orig_stdout, orig_stderr, orig_stdin = sys.stdout, sys.stderr, sys.stdin
        sys.stdout, sys.stderr, sys.stdin = stdout_buf, stderr_buf, io.StringIO(stdin_data or "")

        # Lint for dangerous constructs
        _lint_code_for_danger(code, allow_imports)

        # Prepare restricted globals
        safe_globals = {"__builtins__": SAFE_BUILTINS.copy()}
        if not allow_imports:
            safe_globals["__builtins__"]["__import__"] = _restricted_import

        # Execute user code
        exec(compile(code, filename="<user_code>", mode="exec"), safe_globals, None)

    except Exception:
        traceback.print_exc(file=stderr_buf)
    finally:
        # Restore std streams
        try:
            sys.stdout, sys.stderr, sys.stdin = orig_stdout, orig_stderr, orig_stdin
        except Exception:
            pass
        elapsed = time.time() - start
        q.put({
            "stdout": stdout_buf.getvalue(),
            "stderr": stderr_buf.getvalue(),
            "elapsed": elapsed,
        })


def run_code(code: str, stdin_data: str = "", timeout_s: float = 3.0, allow_imports: bool = False, mem_limit_mb: int = 256):
    """Execute code in a separate process with time & memory limits."""
    q: Queue = Queue()
    p = Process(target=_worker, args=(code, stdin_data, allow_imports, int(mem_limit_mb), int(max(1, timeout_s)), q))
    p.start()
    p.join(timeout=timeout_s)

    if p.is_alive():
        p.terminate()
        try:
            p.join(1)
        except Exception:
            pass
        return "", "Execution timed out.",  timeout_s

    try:
        result = q.get_nowait()
    except Exception:
        return "", "No output captured (possibly killed).", 0.0

    return result.get("stdout", ""), result.get("stderr", ""), float(result.get("elapsed", 0.0))


# -----------------------------
# Gradio UI
# -----------------------------
EXAMPLE_CODE = """
# Example: Fibonacci numbers under 200
fib = [0, 1]
while fib[-1] + fib[-2] < 200:
    fib.append(fib[-1] + fib[-2])
print("Fibonacci:", fib)
""".strip()


def interface_run(code, stdin_text, timeout_s, allow_imports, mem_limit):
    stdout, stderr, elapsed = run_code(code or "", stdin_text or "", timeout_s or 3.0, allow_imports, mem_limit or 256)
    status = f"Finished in {elapsed:.3f}s"
    return stdout, stderr, status


with gr.Blocks(title="Python Code Runner (Sandbox)") as demo:
    gr.Markdown("""
    # 🐍 Python Code Runner (Sandbox)
    A lightweight sandbox to test Python snippets inside your Hugging Face Space.

    **Notes**
    - Imports are **disabled by default** for safety. You can enable them (unsafe) via the checkbox.
    - CPU and memory limits help protect the Space from heavy or infinite computations.
    - `input()` is disabled; use the *stdin* box instead.
    """)

    with gr.Row():
        code = gr.Code(language="python", value=EXAMPLE_CODE, lines=18, label="Python code")
    with gr.Row():
        stdin_tb = gr.Textbox(lines=3, label="stdin (optional)")
    with gr.Row():
        timeout = gr.Slider(1, 20, value=3, step=1, label="Time limit (seconds)")
        mem_limit = gr.Slider(64, 2048, value=256, step=64, label="Memory limit (MB)")
        allow_imports = gr.Checkbox(False, label="Allow imports (unsafe)")
    run_btn = gr.Button("Run", variant="primary")

    with gr.Row():
        stdout = gr.Textbox(lines=12, label="stdout", interactive=False)
    with gr.Row():
        stderr = gr.Textbox(lines=8, label="stderr", interactive=False)
    status = gr.Label(value="Idle", label="Status")

    run_btn.click(interface_run, inputs=[code, stdin_tb, timeout, allow_imports, mem_limit], outputs=[stdout, stderr, status])

    gr.Examples(
        examples=[
            ["print('Hello from Space!')"],
            ["for i in range(5):\n    print(i*i)"] ,
            ["# Using stdin:\nname = input('Name: ')\nprint('Hi', name)"]
        ],
        inputs=[code],
        label="Quick examples"
    )

if __name__ == "__main__":
    demo.launch()