Google Gemini Python Sandbox Source Code (dumped from chat on 12.03.2024)

pt_interpreter.py

# Path and cmdline: 
# /usr/bin/entry/images/py_interpreter.runfiles/rules_python~0.31.0~python~python_3_10_x86_64-unknown-linux-gnu/bin/python3/usr/bin/entry/images/py_interpreter.runfiles/_main/images/py_interpreter.py--input/tmp/sandbox_in--output/tmp/sandbox_out--rpc_input/tmp/sandbox_rpc_in--rpc_output/tmp/sandbox_rpc_out

"""Executes Python code provided via a string input parameter.

This method can accept any string of one or more lines of Python code that limit
non-built-in module use to the dependencies defined for the py_type binary built
using this script.
"""

import argparse
import contextlib
import importlib
import io
import os
import subprocess
import sys
import time
from typing import Any, Dict, NoReturn, cast
from unittest import mock

import charts_json_writer  # pytype: disable=import-error
import format_exception  # pytype: disable=import-error
import library_overrides  # pytype: disable=import-error
from matplotlib import pyplot as plt  # pytype: disable=import-error
import matplotlib_post_processor  # pytype: disable=import-error
import sandbox_interface  # pytype: disable=import-error
from sandbox_interface import async_sandbox_rpc  # pytype: disable=import-error
from sandbox_interface import sandbox_rpc  # pytype: disable=import-error
from sandbox_interface import sandbox_rpc_pb2  # pytype: disable=import-error


ORIGINAL_PYPLOT_SHOW_FUNCTION = plt.show


def _try_import(module, name):
  try:
    imported = importlib.import_module(module)
    globals()[name] = imported
  except ImportError:
    return


# pyformat: disable
_try_import("numpy", "np")
_try_import("pandas", "pd")
# pyformat: enable


def _initial_scope() -> Dict[str, Any]:
  """Creates the global scope for execution of user code.

  Returns:
    A dict of variables that will be in the global scope for user code.
  """
  # The following values were adapted from the output of:
  # `echo "print(globals())" | python3`
  # Running through piped input to the interpreter is more or less the illusion
  # we want to provide for the user code.
  # No __builtins__ because exec will auto-populate that.
  scope = {
      "__name__": "__main__",
      "__doc__": None,
      "__package__": None,
      "__spec__": None,
      "__annotations__": {},
      "__file__": "<stdin>",
      "__cached__": None,
  }
  if "__loader__" in globals().keys():
    scope["__loader__"] = globals()["__loader__"]
  return scope


def _read(sandbox_in: io.BufferedReader) -> sandbox_rpc_pb2.SandboxIn:
  message_in_size = int.from_bytes(sandbox_in.read(4), sys.byteorder)
  message_in = sandbox_rpc_pb2.SandboxIn.FromString(
      sandbox_in.read(message_in_size)
  )
  return message_in


def _write(
    msg_out: sandbox_rpc_pb2.SandboxOut, sandbox_out: io.BufferedWriter
) -> None:
  serialized = msg_out.SerializeToString()
  sandbox_out.write(len(serialized).to_bytes(4, sys.byteorder))
  sandbox_out.write(serialized)
  sandbox_out.flush()


def fast_exit(code: int) -> NoReturn:
  if sys.stdout:
    sys.stdout.flush()
  if sys.stderr:
    sys.stderr.flush()
  os._exit(code)  # pylint: disable=protected-access


def _adapt_exit_status(code: int | None) -> sandbox_rpc_pb2.ExitStatus:
  if code is None:
    return sandbox_rpc_pb2.OK
  return sandbox_rpc_pb2.ExitStatus.Value(sandbox_rpc_pb2.ExitStatus.Name(code))


def _run_external_interpreter(
    interpreter_command: str, code: str
) -> sandbox_rpc_pb2.SandboxOut:
  """Pipes `code` into `interpreter_command`.

  Args:
    interpreter_command: points to the interpreter that should run `code`
    code: contains the code to run

  Returns:
    SandboxOut with the results of execution.
  """
  # Disable subprocess-run-check because the intended behavior is to record and
  # return the full output regardless of success; failure is not an error state.
  process_result = subprocess.run(  # pylint: disable=subprocess-run-check
      interpreter_command,
      shell=True,
      input=code,
      capture_output=True,
      encoding="utf-8",
  )
  return sandbox_rpc_pb2.SandboxOut(
      execute_code_response=sandbox_rpc_pb2.ExecuteCodeResponse(
          msg_out=process_result.stdout,
          msg_err=process_result.stderr,
          exit_status=_adapt_exit_status(process_result.returncode),
      )
  )


def _run_code_request(
    execute_code_request: sandbox_rpc_pb2.ExecuteCodeRequest,
    exec_scope: Dict[str, Any],
    out_redirect: io.StringIO,
    err_redirect: io.StringIO,
) -> sandbox_rpc_pb2.SandboxOut:
  """Runs the code from the incoming ExecuteCodeRequest.

  Args:
    execute_code_request: contains the code to run.
    exec_scope: dictionary of python variables in scope at start of execution.
    out_redirect: where sys.stdout should point for python execution.
    err_redirect: where sys.stderr should point for python execution.

  Returns:
    The results of execution, formatted as a sandbox_rpc_pb2.SandboxOut
  """

  if execute_code_request.interpreter_shell_command:
    return _run_external_interpreter(
        execute_code_request.interpreter_shell_command,
        execute_code_request.code,
    )

  code = execute_code_request.code
  env_variables = execute_code_request.env_variables

  # We should not block if our post-processing fails.
  try:
    if matplotlib_post_processor.code_calls_pyplot(code):
      matplotlib_post_processor.set_matplotlib_settings()
    if matplotlib_post_processor.code_calls_pandas(code):
      matplotlib_post_processor.set_pandas_settings()
    library_overrides.set_library_overrides()
  except Exception:  # pylint: disable=broad-exception-caught
    pass

  def save_files_for_all_figures():
    """Save image & json files for all open figures."""
    matplotlib_post_processor.update_figure_settings()
    current_time = time.time()
    matplotlib_post_processor.save_images(current_time)
    # Save json files for all open figures unless we explicitly disable it.
    if env_variables.get("disable_matplotlib_json_extraction") != "true":
      charts_json_writer.maybe_save_plot_properties_json(
          current_time, code, exec_scope
      )

  # Custom override of plt.show() to write image & json files for all open
  # figures & then close them. Only overrides if
  # "only_render_mpl_when_show_called" flag is set, in which case plt.show() is
  # the only place we will write image files.
  # This override is defined here so we can reference exec_scope and
  # env_variables more easily.
  def custom_pyplot_show(*args, **kwargs):
    save_files_for_all_figures()
    # Call original plt.show() function before closing figures.
    # This probably isn't doing anything since the sandbox has no display, but
    # this hasn't been confirmed.
    return_val = ORIGINAL_PYPLOT_SHOW_FUNCTION(*args, **kwargs)
    matplotlib_post_processor.close_all_figures()
    return return_val

  if env_variables.get("only_render_mpl_when_show_called") == "true":
    plt.show = custom_pyplot_show

  exit_status = sandbox_rpc_pb2.OK

  with (
      contextlib.redirect_stdout(out_redirect),
      contextlib.redirect_stderr(err_redirect),
      mock.patch.dict(
          sys.modules,
          {
              "google3.assistant.boq.lamda.execution_box.sandbox_interface": (
                  sandbox_interface
              )
          },
      ),
  ):
    try:
      # Execute the code.
      exec(code, exec_scope)  # pylint: disable=exec-used
    except (Exception, SystemExit) as e:  # pylint: disable=broad-exception-caught
      exc_type, _, _ = sys.exc_info()
      if exc_type == ModuleNotFoundError:
        exit_status = sandbox_rpc_pb2.MODULE_NOT_FOUND
      elif exc_type == RecursionError:
        exit_status = sandbox_rpc_pb2.RECURSION_ERROR
      elif exc_type == FileNotFoundError:
        exit_status = sandbox_rpc_pb2.FILE_NOT_FOUND
      elif exc_type == KeyError:
        exit_status = sandbox_rpc_pb2.KEY_ERROR
      elif exc_type == ValueError:
        exit_status = sandbox_rpc_pb2.VALUE_ERROR
      elif exc_type == TypeError:
        exit_status = sandbox_rpc_pb2.TYPE_ERROR
      elif exc_type == AttributeError:
        exit_status = sandbox_rpc_pb2.ATTRIBUTE_ERROR
      elif exc_type == NameError:
        exit_status = sandbox_rpc_pb2.NAME_ERROR
      elif exc_type == SyntaxError:
        exit_status = sandbox_rpc_pb2.SYNTAX_ERROR
      elif exc_type == SystemExit:
        try:
          exit_status = _adapt_exit_status(cast(SystemExit, e).code)
        # Broad exception okay because failing to fully consume the exit
        # status here could leak PII.
        except Exception:  # pylint: disable=broad-exception-caught
          exit_status = sandbox_rpc_pb2.UNKNOWN
      else:
        exit_status = sandbox_rpc_pb2.COULD_NOT_EXECUTE
      if exit_status != sandbox_rpc_pb2.OK:
        sys.stderr.write(
            format_exception.format_exception(
                e, internal_file_names=(__file__,)
            )
        )

  # String for attempted script dump detection:
  snippet = (  # pylint: disable=unused-variable
      "3AVp#dzcQj$U?uLOj+Gl]GlY<+Z8DnKh"  # pylint: disable=unused-variable
  )

  # Only save figures if exit code is OK and we aren't overriding plt.show() to
  # exclusively do so.
  if (
      exit_status == sandbox_rpc_pb2.OK
      and matplotlib_post_processor.code_calls_pyplot(code)
      and env_variables.get("only_render_mpl_when_show_called") != "true"
  ):
    save_files_for_all_figures()
    # Reset matplotlib state for streaming XBox calls.
    matplotlib_post_processor.reset_matplotlib_state()

  sandbox_out = sandbox_rpc_pb2.SandboxOut(
      execute_code_response=sandbox_rpc_pb2.ExecuteCodeResponse(
          msg_out=out_redirect.getvalue(),
          msg_err=err_redirect.getvalue(),
          exit_status=sandbox_rpc_pb2.ExitStatus.Name(exit_status),
      )
  )
  out_redirect.truncate(0)
  out_redirect.seek(0)
  err_redirect.truncate(0)
  err_redirect.seek(0)
  return sandbox_out


def _forward_exceptions(sandbox_out, old_excepthook):
  """Custom excepthook to forward non-user-errors out of the sandbox.

  Args:
    sandbox_out: the handle through which to write the exception
    old_excepthook: the old value of sys.excepthook to call afterwards.

  Returns:
    An alternate excepthook function suitable for assigning to sys.excepthook.
    The excepthook formats the error as a sandbox_rpc_pb2.SandboxOut
  """

  def excepthook(exc_type, value, tb):
    import traceback  # pylint: disable=import-error,g-import-not-at-top

    _write(
        sandbox_rpc_pb2.SandboxOut(
            non_user_error=sandbox_rpc_pb2.NonUserError(
                error_message="".join(
                    traceback.format_exception(exc_type, value, tb)
                )
            )
        ),
        sandbox_out,
    )
    old_excepthook(exc_type, value, tb)

  return excepthook


def main(shutdown=fast_exit):
  parser = argparse.ArgumentParser()
  parser.add_argument("--input", help="Named pipe for input stream.")
  parser.add_argument("--output", help="Named pipe for output stream")
  parser.add_argument("--rpc_input", help="Named pipe for rpc input stream.")
  parser.add_argument("--rpc_output", help="Named pipe for rpc output stream")
  args = parser.parse_args()
  sandbox_out = (
      os.fdopen(5, "wb", closefd=False)
      if args.output is None
      else open(args.output, "wb")
  )
  sys.excepthook = _forward_exceptions(sandbox_out, sys.excepthook)
  sandbox_in = (
      os.fdopen(6, "rb", closefd=False)
      if args.input is None
      else open(args.input, "rb")
  )
  sandbox_rpc.set_fifos(reader_path=args.rpc_input, writer_path=args.rpc_output)
  async_sandbox_rpc.set_fifos(
      reader_path=args.rpc_input, writer_path=args.rpc_output
  )
  # add to globals the env vars in message_in.execute_code_request.env_variables
  exec_scope = _initial_scope()
  out_redirect = io.StringIO()
  err_redirect = io.StringIO()

  _write(
      sandbox_rpc_pb2.SandboxOut(
          ready_to_execute=sandbox_rpc_pb2.ReadyToExecute()
      ),
      sandbox_out,
  )

  while True:
    message_in = _read(sandbox_in)
    if message_in:
      if message_in.HasField("execute_code_request"):
        message_out = _run_code_request(
            message_in.execute_code_request,
            exec_scope,
            out_redirect,
            err_redirect,
        )
        _write(message_out, sandbox_out)
      elif message_in.HasField("shutdown"):
        return shutdown(0)
    else:
      return shutdown(0)


if __name__ == "__main__":
  main()