mirror of
https://codeberg.org/privacy1st/snowflake-stats
synced 2024-12-23 01:56:04 +01:00
66 lines
1.9 KiB
Python
66 lines
1.9 KiB
Python
|
#!/usr/bin/env python3
|
||
|
from typing import List
|
||
|
import sys
|
||
|
import subprocess
|
||
|
import shlex
|
||
|
|
||
|
|
||
|
class NonZeroExitCode(Exception):
|
||
|
def __init__(self, exit_code, stdout, stderr):
|
||
|
super(NonZeroExitCode, self).__init__(f'exit_code: {exit_code}, stdout: {stdout}, stderr: {stderr}')
|
||
|
|
||
|
|
||
|
def wrap_command_with_ssh(command: List[str], ssh_host: str) -> List[str]:
|
||
|
return ['ssh', ssh_host, ' '.join([shlex.quote(arg) for arg in command])]
|
||
|
|
||
|
|
||
|
def execute(command: List[str], ssh_host: str = None) -> None:
|
||
|
"""
|
||
|
https://docs.python.org/3.10/library/subprocess.html#frequently-used-arguments
|
||
|
|
||
|
Executes the given command using stdin, stdout, stderr of the parent process.
|
||
|
|
||
|
The output is not captured.
|
||
|
|
||
|
:raises subprocess.CalledProcessError: In case of non-zero exit code.
|
||
|
"""
|
||
|
if ssh_host is not None:
|
||
|
command = wrap_command_with_ssh(command, ssh_host)
|
||
|
|
||
|
print(f'[DEBUG] execute: {command}')
|
||
|
|
||
|
subprocess.run(
|
||
|
command,
|
||
|
stdin=sys.stdin,
|
||
|
stdout=sys.stdout,
|
||
|
stderr=sys.stderr,
|
||
|
check=True,
|
||
|
)
|
||
|
|
||
|
|
||
|
def execute_and_capture(command: List[str], ssh_host: str = None, file='stdout') -> str:
|
||
|
"""
|
||
|
Executes the given command and captures it stdout.
|
||
|
|
||
|
:return: Stdout of subprocess.
|
||
|
:raises NonZeroExitCode: In case of non-zero exit code. This exception includes the exit_code, stdout and stderr in it's message.
|
||
|
"""
|
||
|
if ssh_host is not None:
|
||
|
command = wrap_command_with_ssh(command, ssh_host)
|
||
|
|
||
|
print(f'[DEBUG] execute_and_capture: {command}')
|
||
|
|
||
|
completed: subprocess.CompletedProcess = subprocess.run(
|
||
|
command,
|
||
|
capture_output=True,
|
||
|
text=True,
|
||
|
)
|
||
|
if completed.returncode != 0:
|
||
|
raise NonZeroExitCode(completed.returncode, completed.stdout, completed.stderr)
|
||
|
|
||
|
if file == 'stdout':
|
||
|
return completed.stdout
|
||
|
if file == 'stderr':
|
||
|
return completed.stderr
|
||
|
raise ValueError('invalid argument')
|