From ede604cfc55a18d5f3a776805a554f126afe1cc3 Mon Sep 17 00:00:00 2001
From: Austin Bozowski <bozowski@google.com>
Date: Wed, 1 Jun 2022 07:35:53 +0000
Subject: [PATCH] Merge from pull/19006

---
 examples/chef/stateful_shell.py | 42 +++++++++++++++++++++++++++------
 1 file changed, 35 insertions(+), 7 deletions(-)

diff --git a/examples/chef/stateful_shell.py b/examples/chef/stateful_shell.py
index 56ba52bf493464..20615796c88ae4 100644
--- a/examples/chef/stateful_shell.py
+++ b/examples/chef/stateful_shell.py
@@ -16,11 +16,13 @@
 import shlex
 import subprocess
 import sys
-import tempfile
-from typing import Dict
+from typing import Dict, Optional
 
 import constants
 
+_ENV_FILENAME = ".shell_env"
+_OUTPUT_FILENAME = ".shell_output"
+_HERE = os.path.dirname(os.path.abspath(__file__))
 
 TermColors = constants.TermColors
 
@@ -42,7 +44,8 @@ def __init__(self) -> None:
 
         # This file holds the env after running a command. This is a better approach
         # than writing to stdout because commands could redirect the stdout.
-        self.envfile_path: str = os.path.join(tempfile.gettempdir(), "envfile")
+        self.envfile_path: str = os.path.join(_HERE, _ENV_FILENAME)
+        self.cmd_output_path: str = os.path.join(_HERE, _OUTPUT_FILENAME)
 
     def print_env(self) -> None:
         """Print environment variables in commandline friendly format for export.
@@ -56,21 +59,41 @@ def print_env(self) -> None:
             if env_var:
                 print(f"export {env_var}={quoted_value}")
 
-    def run_cmd(self, cmd: str, *, raise_on_returncode=False) -> None:
+    def run_cmd(
+        self, cmd: str, *,
+        raise_on_returncode=False,
+        return_cmd_output=False,
+    ) -> Optional[str]:
         """Runs a command and updates environment.
 
         Args:
           cmd: Command to execute.
+            This does not support commands that run in the background e.g. `<cmd> &`
           raise_on_returncode: Whether to raise an error if the return code is nonzero.
+          return_cmd_output: Whether to return the command output.
+            If enabled, the text piped to screen won't be colorized due to output
+            being passed through `tee`.
 
         Raises:
           RuntimeError: If raise_on_returncode is set and nonzero return code is given.
+
+        Returns:
+          Output of command if return_cmd_output set to True.
         """
         env_dict = {}
-
         # Set OLDPWD at beginning because opening the shell clears this. This handles 'cd -'.
         # env -0 prints the env variables separated by null characters for easy parsing.
-        command_with_state = f"OLDPWD={self.env.get('OLDPWD', '')}; {cmd}; env -0 > {self.envfile_path}"
+
+        if return_cmd_output:
+            # Piping won't work here because piping will affect how environment variables
+            # are propagated. This solution uses tee without piping to preserve env variables.
+            redirect = f" > >(tee \"{self.cmd_output_path}\") 2>&1 "  # include stderr
+        else:
+            redirect = ""
+
+        command_with_state = (
+            f"OLDPWD={self.env.get('OLDPWD', '')}; {cmd} {redirect};"
+            f" env -0 > {self.envfile_path}")
         with subprocess.Popen(
             [command_with_state],
             env=self.env, cwd=self.cwd,
@@ -79,7 +102,7 @@ def run_cmd(self, cmd: str, *, raise_on_returncode=False) -> None:
             returncode = proc.wait()
 
         # Load env state from envfile.
-        with open(self.envfile_path) as f:
+        with open(self.envfile_path, encoding="latin1") as f:
             # Split on null char because we use env -0.
             env_entries = f.read().split("\0")
             for entry in env_entries:
@@ -92,3 +115,8 @@ def run_cmd(self, cmd: str, *, raise_on_returncode=False) -> None:
         if raise_on_returncode and returncode != 0:
             raise RuntimeError(
                 f"Error. Return code is not 0. It is: {returncode}")
+
+        if return_cmd_output:
+            with open(self.cmd_output_path, encoding="latin1") as f:
+                output = f.read()
+            return output