astoeckel/sandbox.py

## sandbox.py
#!/usr/bin/env python3

"""
This module provides an `eval_in_sandbox` function that, on Linux systems,
executes arbitrary Python code in a secure and lightweight sandbox with
memory, I/O, and time limits.

----

Copyright (c) 2023 Andreas Stöckel

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import io
import os
import shutil
import subprocess
import tempfile
import time
import sys
import dataclasses
import select


@dataclasses.dataclass
class ProcessOutput:
    returncode: int = 0
    stdout: str = ""
    stderr: str = ""

    def __bool__(self):
        return self.returncode == 0


class SandboxError(Exception):
    pass


def eval_in_sandbox(
        python_code: str,
        mem_max_mb: int = 100,
        cpu_max_time_sec: int = 1,
        timeout_max_time_sec: int = 10,
        io_max_kb: int = 128,
) -> ProcessOutput:
    """
    Executes the given script in an isolated sandbox with memory, CPU, and I/O
    limits. Returns the captured stdout and stderr.  Throws `SandboxError` if
    anything goes wrong.

    :param python_code: is a string containing the Python code that should be
        safely executed.
    :param mem_max_mb: is the maximum size of the address space used by the
        child process in mebibytes. Note that this value must be substantially
        larger than the actual amount of memory used by the process, since it
        includes the memory needed by the Python interpreter executable itself.
    :param cpu_max_time_sec: is the maximum CPU time that can be used by the
        child process. CPU time is the time during which the process is
        actually active; this doesn't include sleeping, or waiting for I/O.
    :param timeout_max_time_sec: is the maximum wall-clock time for which the
        process may be alive.
    :param io_max_kb: maximum number of kibibytes that may be written by the
        process to stdout/stdin.
    :return: an instance of the `ProcessOutput` structure, containing the
        process return code and stdin/stdout.
    """

    # Convert the given memory size limits to bytes
    mem_max_bytes = 1024 * 1024 * mem_max_mb
    io_max_bytes = 1024 * io_max_kb

    # Search for the bubblewrap executable
    bwrap_exe = shutil.which("bwrap")
    if bwrap_exe is None:
        raise SandboxError("Cannot find the `bwrap` executable")

    # Search for the prlimit executable
    prlimit_exe = shutil.which("prlimit")
    if prlimit_exe is None:
        raise SandboxError("Cannot find the `prlimit` executable")

    # Determine the path to the Python interpreter; sometimes `sys.executable`
    # returns wrong results, so we use `shutil.which` as a fallback
    python_exe = sys.executable
    if not os.access(python_exe, os.X_OK):
        python_exe = shutil.which("python3")
        if python_exe is None:
            raise SandboxError("Cannot find the Python executable")
    if not python_exe.startswith("/usr"):
        raise SandboxError("Python interpreter not located in `/usr`")

    # Securely create a temporary directory; this directory will automatically
    # be deleted
    with tempfile.TemporaryDirectory() as tmp_dir:
        # Write the given code into the temporary directory
        code_filename = os.path.join(tmp_dir, "code.py")
        with open(code_filename, "w", encoding="utf-8") as f:
            f.write(python_code)
            f.write("\n")

        # Assemble the arguments that we need to pass to `bwrap`
        args = [
            # Use `prlimit` to limit the CPU time and the size of the address
            # space used by the child commands
            prlimit_exe,

            # Set the soft and hard cpu time limit
            f"--cpu={cpu_max_time_sec}:{cpu_max_time_sec}",

            # Set the soft and hard address space limit
            f"--as={mem_max_bytes}:{mem_max_bytes}",

            # Execute bwrap
            bwrap_exe,

            # Use an isolated terminal session for the subprcoess
            "--new-session",

            # Isolate all possible Linux namespaces.
            "--unshare-all",

            # Clear all environment variables
            "--clearenv",

            # Execute the script inside the "home" directory
            "--chdir", "/home",

            # Mount the temporary directory as home
            "--ro-bind", tmp_dir, "/home",

            # Make the directory containing the Python interpreter available
            "--ro-bind", "/usr", "/usr",
            "--ro-bind", "/lib", "/lib",
            "--ro-bind", "/lib64", "/lib64",

            # Lastly, make the virtual root directory read-only
            "--remount-ro", "/",

            # Execute the Python interpreter
            python_exe,

            # Disable buffering
            "-u",

            # Execute the script
            "/home/code.py"
        ]

        # Execute the subprocess; read from stdout/stderr while handling
        # timeouts and I/O limits.
        try:
            # Create the subprocess; close stdin
            proc = subprocess.Popen(
                args,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                stdin=subprocess.DEVNULL
            )

            # Mark stdout and stderr as non-blocking; this is needed because we
            # always try to read 1024-byte chunks from these streams.
            os.set_blocking(proc.stdout.fileno(), False)
            os.set_blocking(proc.stderr.fileno(), False)

            # Get a second timestamp from a monotonic time source
            t_start = time.monotonic()

            # Read stdout and stderr
            buf_stdout, buf_stderr = io.BytesIO(), io.BytesIO()
            src_fds = [proc.stdout, proc.stderr]
            tar_bufs = [buf_stdout, buf_stderr]
            returncode = None
            while returncode is None:
                # Abort if we get a timeout
                t_now = time.monotonic()
                if t_now - t_start > timeout_max_time_sec:
                    raise SandboxError("Process timed out")

                # Wait for data being available on stdout/stderr
                r_fds, _, x_fds = select.select(
                    src_fds,

                    # We do not have any input stream
                    [],

                    # Wait for errors/the proc being closed
                    src_fds,

                    # Wait for at most 100ms
                    0.1,
                )

                # Read data from stdout/stderr; abort if the I/O buffers are
                # too large
                for src_fd, tar_buf in list(zip(src_fds, tar_bufs)):
                    if src_fd in r_fds:
                        # FD is ready for reading, try to read a bit of data
                        buf = src_fd.read(1024)

                        # Did that fail? If yes, do not bother asking select
                        # for that fd anymore
                        if not buf:
                            src_fds.remove(src_fd)
                            tar_bufs.remove(tar_buf)
                            continue

                        # Append the data we read to the target buffer
                        tar_buf.write(buf)
                        if tar_buf.tell() > io_max_bytes:
                            raise SandboxError("I/O buffer size exceeded")

                # Discard stdout/stderr if closed by the process
                for src_fd, tar_buf in list(zip(src_fds, tar_bufs)):
                    if src_fd in x_fds:
                        # Do not bother about this FD anymore
                        src_fds.remove(src_fd)
                        tar_bufs.remove(tar_buf)

                # Wait for the process to exit, if stdout/stderr were closed
                if not src_fds:
                    try:
                        returncode = proc.wait(0.1)
                    except TimeoutError:
                        # If we get a timeout, continue in the loop; our own
                        # high-level timeout code will trigger at some point.
                        pass

            # Try to decode the stdout/stderr buffer
            try:
                stdout = str(buf_stdout.getvalue(), "utf-8")
                stderr = str(buf_stderr.getvalue(), "utf-8")
            except UnicodeDecodeError:
                raise SandboxError("Subprocess did not produce valid UTF-8")

            # Strip spaces away
            stdout = stdout.strip()
            stderr = stderr.strip()

            # Return the output!
            return ProcessOutput(
                returncode=returncode,
                stderr=stderr,
                stdout=stdout,
            )
        except SandboxError:
            # Kill the subprocess
            if returncode is None:
                proc.kill()
                proc.wait()

            # Re-raise the exception
            raise


if __name__ == "__main__":
    def main():
        """
        Small test program.
        """
        res = eval_in_sandbox("""
print('Hello World from stdout!')

import os
import json
print("My UID is", os.getuid())
print("My GID is", os.getgid())
print("My PID is", os.getpid())
print("My CWD is", os.getcwd())
print("My environment is", json.dumps(dict(**os.environ), indent=4))

try:
    with open("./foo.txt", "w") as f:
        f.write("EVIL\\n")
except BaseException as e:
    print(f"Writing to a file failed with: {e}")

try:
    os.unlink("code.py")
except BaseException as e:
    print(f"Deleting a file failed with: {e}")


import sys
sys.stderr.write('Hello World from stderr!')

sys.exit(42)
""")
        print(f"Stdout:     {res.stdout}")
        print(f"Stderr:     {res.stderr}")
        print(f"Returncode: {res.returncode}")
        sys.exit(res.returncode)


    main()
	#!/usr/bin/env python3

	"""
	This module provides an `eval_in_sandbox` function that, on Linux systems,
	executes arbitrary Python code in a secure and lightweight sandbox with
	memory, I/O, and time limits.

	----

	Copyright (c) 2023 Andreas Stöckel

	Permission is hereby granted, free of charge, to any person obtaining a copy
	of this software and associated documentation files (the "Software"), to deal
	in the Software without restriction, including without limitation the rights
	to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	copies of the Software, and to permit persons to whom the Software is
	furnished to do so, subject to the following conditions:

	The above copyright notice and this permission notice shall be included in all
	copies or substantial portions of the Software.

	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
	SOFTWARE.
	"""

	import io
	import os
	import shutil
	import subprocess
	import tempfile
	import time
	import sys
	import dataclasses
	import select


	@dataclasses.dataclass
	class ProcessOutput:
	returncode: int = 0
	stdout: str = ""
	stderr: str = ""

	def __bool__(self):
	return self.returncode == 0


	class SandboxError(Exception):
	pass


	def eval_in_sandbox(
	python_code: str,
	mem_max_mb: int = 100,
	cpu_max_time_sec: int = 1,
	timeout_max_time_sec: int = 10,
	io_max_kb: int = 128,
	) -> ProcessOutput:
	"""
	Executes the given script in an isolated sandbox with memory, CPU, and I/O
	limits. Returns the captured stdout and stderr. Throws `SandboxError` if
	anything goes wrong.

	:param python_code: is a string containing the Python code that should be
	safely executed.
	:param mem_max_mb: is the maximum size of the address space used by the
	child process in mebibytes. Note that this value must be substantially
	larger than the actual amount of memory used by the process, since it
	includes the memory needed by the Python interpreter executable itself.
	:param cpu_max_time_sec: is the maximum CPU time that can be used by the
	child process. CPU time is the time during which the process is
	actually active; this doesn't include sleeping, or waiting for I/O.
	:param timeout_max_time_sec: is the maximum wall-clock time for which the
	process may be alive.
	:param io_max_kb: maximum number of kibibytes that may be written by the
	process to stdout/stdin.
	:return: an instance of the `ProcessOutput` structure, containing the
	process return code and stdin/stdout.
	"""

	# Convert the given memory size limits to bytes
	mem_max_bytes = 1024 * 1024 * mem_max_mb
	io_max_bytes = 1024 * io_max_kb

	# Search for the bubblewrap executable
	bwrap_exe = shutil.which("bwrap")
	if bwrap_exe is None:
	raise SandboxError("Cannot find the `bwrap` executable")

	# Search for the prlimit executable
	prlimit_exe = shutil.which("prlimit")
	if prlimit_exe is None:
	raise SandboxError("Cannot find the `prlimit` executable")

	# Determine the path to the Python interpreter; sometimes `sys.executable`
	# returns wrong results, so we use `shutil.which` as a fallback
	python_exe = sys.executable
	if not os.access(python_exe, os.X_OK):
	python_exe = shutil.which("python3")
	if python_exe is None:
	raise SandboxError("Cannot find the Python executable")
	if not python_exe.startswith("/usr"):
	raise SandboxError("Python interpreter not located in `/usr`")

	# Securely create a temporary directory; this directory will automatically
	# be deleted
	with tempfile.TemporaryDirectory() as tmp_dir:
	# Write the given code into the temporary directory
	code_filename = os.path.join(tmp_dir, "code.py")
	with open(code_filename, "w", encoding="utf-8") as f:
	f.write(python_code)
	f.write("\n")

	# Assemble the arguments that we need to pass to `bwrap`
	args = [
	# Use `prlimit` to limit the CPU time and the size of the address
	# space used by the child commands
	prlimit_exe,

	# Set the soft and hard cpu time limit
	f"--cpu={cpu_max_time_sec}:{cpu_max_time_sec}",

	# Set the soft and hard address space limit
	f"--as={mem_max_bytes}:{mem_max_bytes}",

	# Execute bwrap
	bwrap_exe,

	# Use an isolated terminal session for the subprcoess
	"--new-session",

	# Isolate all possible Linux namespaces.
	"--unshare-all",

	# Clear all environment variables
	"--clearenv",

	# Execute the script inside the "home" directory
	"--chdir", "/home",

	# Mount the temporary directory as home
	"--ro-bind", tmp_dir, "/home",

	# Make the directory containing the Python interpreter available
	"--ro-bind", "/usr", "/usr",
	"--ro-bind", "/lib", "/lib",
	"--ro-bind", "/lib64", "/lib64",

	# Lastly, make the virtual root directory read-only
	"--remount-ro", "/",

	# Execute the Python interpreter
	python_exe,

	# Disable buffering
	"-u",

	# Execute the script
	"/home/code.py"
	]

	# Execute the subprocess; read from stdout/stderr while handling
	# timeouts and I/O limits.
	try:
	# Create the subprocess; close stdin
	proc = subprocess.Popen(
	args,
	stdout=subprocess.PIPE,
	stderr=subprocess.PIPE,
	stdin=subprocess.DEVNULL
	)

	# Mark stdout and stderr as non-blocking; this is needed because we
	# always try to read 1024-byte chunks from these streams.
	os.set_blocking(proc.stdout.fileno(), False)
	os.set_blocking(proc.stderr.fileno(), False)

	# Get a second timestamp from a monotonic time source
	t_start = time.monotonic()

	# Read stdout and stderr
	buf_stdout, buf_stderr = io.BytesIO(), io.BytesIO()
	src_fds = [proc.stdout, proc.stderr]
	tar_bufs = [buf_stdout, buf_stderr]
	returncode = None
	while returncode is None:
	# Abort if we get a timeout
	t_now = time.monotonic()
	if t_now - t_start > timeout_max_time_sec:
	raise SandboxError("Process timed out")

	# Wait for data being available on stdout/stderr
	r_fds, _, x_fds = select.select(
	src_fds,

	# We do not have any input stream
	[],

	# Wait for errors/the proc being closed
	src_fds,

	# Wait for at most 100ms
	0.1,
	)

	# Read data from stdout/stderr; abort if the I/O buffers are
	# too large
	for src_fd, tar_buf in list(zip(src_fds, tar_bufs)):
	if src_fd in r_fds:
	# FD is ready for reading, try to read a bit of data
	buf = src_fd.read(1024)

	# Did that fail? If yes, do not bother asking select
	# for that fd anymore
	if not buf:
	src_fds.remove(src_fd)
	tar_bufs.remove(tar_buf)
	continue

	# Append the data we read to the target buffer
	tar_buf.write(buf)
	if tar_buf.tell() > io_max_bytes:
	raise SandboxError("I/O buffer size exceeded")

	# Discard stdout/stderr if closed by the process
	for src_fd, tar_buf in list(zip(src_fds, tar_bufs)):
	if src_fd in x_fds:
	# Do not bother about this FD anymore
	src_fds.remove(src_fd)
	tar_bufs.remove(tar_buf)

	# Wait for the process to exit, if stdout/stderr were closed
	if not src_fds:
	try:
	returncode = proc.wait(0.1)
	except TimeoutError:
	# If we get a timeout, continue in the loop; our own
	# high-level timeout code will trigger at some point.
	pass

	# Try to decode the stdout/stderr buffer
	try:
	stdout = str(buf_stdout.getvalue(), "utf-8")
	stderr = str(buf_stderr.getvalue(), "utf-8")
	except UnicodeDecodeError:
	raise SandboxError("Subprocess did not produce valid UTF-8")

	# Strip spaces away
	stdout = stdout.strip()
	stderr = stderr.strip()

	# Return the output!
	return ProcessOutput(
	returncode=returncode,
	stderr=stderr,
	stdout=stdout,
	)
	except SandboxError:
	# Kill the subprocess
	if returncode is None:
	proc.kill()
	proc.wait()

	# Re-raise the exception
	raise


	if __name__ == "__main__":
	def main():
	"""
	Small test program.
	"""
	res = eval_in_sandbox("""
	print('Hello World from stdout!')

	import os
	import json
	print("My UID is", os.getuid())
	print("My GID is", os.getgid())
	print("My PID is", os.getpid())
	print("My CWD is", os.getcwd())
	print("My environment is", json.dumps(dict(**os.environ), indent=4))

	try:
	with open("./foo.txt", "w") as f:
	f.write("EVIL\\n")
	except BaseException as e:
	print(f"Writing to a file failed with: {e}")

	try:
	os.unlink("code.py")
	except BaseException as e:
	print(f"Deleting a file failed with: {e}")


	import sys
	sys.stderr.write('Hello World from stderr!')

	sys.exit(42)
	""")
	print(f"Stdout: {res.stdout}")
	print(f"Stderr: {res.stderr}")
	print(f"Returncode: {res.returncode}")
	sys.exit(res.returncode)


	main()