admin_rb
107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
import os
|
|
import shutil
|
|
import signal
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import Sequence
|
|
|
|
from aim.exceptions import ExternalCommandFailed, OperationCancelled
|
|
|
|
_presenter = ContextVar('aim_external_presenter', default=None)
|
|
|
|
|
|
@contextmanager
|
|
def external_presentation(presenter):
|
|
"""Install presentation per application context, not in business managers."""
|
|
token = _presenter.set(presenter)
|
|
try:
|
|
yield
|
|
finally:
|
|
_presenter.reset(token)
|
|
|
|
|
|
def stable_cwd(preferred: Path | None = None) -> Path:
|
|
candidates = [preferred, Path('/etc/ansible'), Path.home(), Path('/')]
|
|
for candidate in candidates:
|
|
if candidate is None:
|
|
continue
|
|
try:
|
|
if candidate.is_dir():
|
|
return candidate
|
|
except OSError:
|
|
pass
|
|
return Path('/')
|
|
|
|
|
|
def command_available(command: str) -> bool:
|
|
return shutil.which(command) is not None
|
|
|
|
|
|
def stop_process(process: subprocess.Popen, *, group: bool = False) -> None:
|
|
"""Reap a cancelled command; do not claim to undo work on remote systems."""
|
|
if process.poll() is not None:
|
|
return
|
|
try:
|
|
if group and os.name == 'posix':
|
|
os.killpg(process.pid, signal.SIGTERM)
|
|
else:
|
|
process.terminate()
|
|
process.wait(timeout=3)
|
|
except subprocess.TimeoutExpired:
|
|
if group and os.name == 'posix':
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
else:
|
|
process.kill()
|
|
process.wait()
|
|
except ProcessLookupError:
|
|
process.wait()
|
|
|
|
|
|
def execute(args, kwargs):
|
|
"""Default terminal behavior, including native password/editor interaction."""
|
|
process = subprocess.Popen(list(args), **kwargs)
|
|
try:
|
|
stdout, stderr = process.communicate()
|
|
return subprocess.CompletedProcess(list(args), process.returncode, stdout, stderr)
|
|
finally:
|
|
# Runs on KeyboardInterrupt, SIGTERM/SystemExit, or a regular exception.
|
|
stop_process(process)
|
|
|
|
|
|
def run_external(args: Sequence[str], *, cwd: Path | None = None,
|
|
capture_output: bool = False, suppress_stdout: bool = False,
|
|
capture_stdout: bool = False, capture_stderr: bool = False,
|
|
text: bool = True, env: dict[str, str] | None = None,
|
|
ui_mode: str = 'auto') -> subprocess.CompletedProcess:
|
|
if (capture_output or capture_stdout) and suppress_stdout:
|
|
raise ValueError('capture_output and suppress_stdout are mutually exclusive')
|
|
if not args:
|
|
raise ValueError('An external command is required')
|
|
kwargs = {'cwd': stable_cwd(cwd), 'text': text}
|
|
if env is not None:
|
|
kwargs['env'] = env
|
|
if capture_output:
|
|
kwargs.update(stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
elif capture_stdout:
|
|
kwargs['stdout'] = subprocess.PIPE
|
|
elif suppress_stdout:
|
|
kwargs['stdout'] = subprocess.DEVNULL
|
|
if capture_stderr:
|
|
kwargs['stderr'] = subprocess.PIPE
|
|
try:
|
|
presenter = _presenter.get()
|
|
if presenter:
|
|
return presenter(list(args), kwargs, capture_output=capture_output or capture_stdout or capture_stderr,
|
|
suppress_stdout=suppress_stdout, ui_mode=ui_mode)
|
|
return execute(args, kwargs)
|
|
except KeyboardInterrupt:
|
|
raise OperationCancelled('External command interrupted. Completed remote work is not rolled back.') from None
|
|
except FileNotFoundError as exc:
|
|
raise ExternalCommandFailed(f'Required command not found: {args[0]}') from exc
|
|
except OSError as exc:
|
|
raise ExternalCommandFailed(f'Could not execute {args[0]}: {exc}') from exc
|