custom-hashicorp-vault/setup-scripts/CommandRunner.py
2025-05-07 06:30:09 -05:00

93 lines
No EOL
4.7 KiB
Python

import os, sys, subprocess, select
from typing import List, Tuple
class CommandRunner:
@staticmethod
def run_command(command: str | List[str], check=True) -> Tuple[int, str, str]:
"""Run a command on the system and return the output
Args:
command (str | List[str]): The command to run
check (bool, optional): If the command should raise an exception if it fails (`check=` for `subprocess.run`). Defaults to True.
Returns:
Tuple[int, str, str]: The return code of the command, the output of the command (on standard out), and the error output of the command (on standard error)
"""
# Copy of the environment variables
new_env = os.environ.copy()
if not check:
# If the check isn't set, leave it out of the `subprocess.run`` call
# This means no exception will be raised if the command fails
result = subprocess.run(command, env=new_env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
return result.returncode, result.stdout.decode('utf-8').strip(), result.stderr.decode('utf-8')
else:
# Run the command and raise an exception if it fails
try:
result = subprocess.run(command, env=new_env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, check=True)
return result.returncode, result.stdout.decode('utf-8').strip(), result.stderr.decode('utf-8').strip()
except subprocess.CalledProcessError as e:
# Log the error and re-raise or handle as appropriate
raise RuntimeError(f"Command '{command}' failed with error: {e.stderr}") from e
@staticmethod
def run_command_in_real_time(command: str | List[str]) -> Tuple[int, str, str]:
"""Similar to `run_command` but prints the output of the command in real-time
Args:
command (str | List[str]): The command to run
Returns:
tuple[int, str, str]: The return code of the command, the output of the command (on standard out), and the error output of the command (on standard error)
"""
# Variables to store the output of the command
stdout_lines: list[str] = []
stderr_lines: list[str] = []
# Copy of the environment variables
my_env = os.environ.copy()
# Run the command
# We use a `with` statement to ensure the process is closed properly
with subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=my_env, shell=True) as proc:
# Call `poll()` initially because we want to emulate a do-while loop
return_code = proc.poll()
# Loop until the process is finished
# In theory this will loop about every second (timeout on the select call)
# Or slightly faster if the process outputs something
while(return_code is None):
# Because `readline` block until it gets a line,
# But the executed command ISN'T guaranteed to output a line every time
# We use `select` to check if there's something to read
# It Waits for 1 second or for the process to output something
rlist, _, _ = select.select([proc.stdout.fileno(), proc.stderr.fileno()], [], [], 1)
# There was something to read from the process's stdout
if proc.stdout.fileno() in rlist:
# Read the line from the process
stdout_line = proc.stdout.readline()
# Add the line to the cumulative output
stdout_lines.append(stdout_line.decode('utf-8').strip())
# Print the output in real-time
print(stdout_line.decode('utf-8').strip())
# There was something to read from the process's stderr
if proc.stderr.fileno() in rlist:
# Read the line from the process
stderr_line = proc.stderr.readline()
# Add the line to the cumulative output
stderr_lines.append(stderr_line.decode('utf-8').strip())
# Print the error output of the command in real-time to stderr
print(stderr_line.decode('utf-8').strip(), file=sys.stderr)
# Update the return code (to see if the process is finished)
return_code = proc.poll()
# Return the error code AND the full output of the command as a string
return return_code, '\n'.join(stdout_lines), '\n'.join(stderr_lines)