# linux shell command execution utility # # Copyright (C) 2020 Eric Callahan # # This file may be distributed under the terms of the GNU GPLv3 license. from __future__ import annotations import os import shlex import logging import signal import asyncio from tornado import gen from utils import ServerError # Annotation imports from typing import ( TYPE_CHECKING, Optional, Callable, Coroutine, Dict, ) if TYPE_CHECKING: from confighelper import ConfigHelper from asyncio import BaseTransport OutputCallback = Optional[Callable[[bytes], None]] class ShellCommandError(ServerError): def __init__(self, message: str, return_code: Optional[int], stdout: Optional[bytes] = b"", stderr: Optional[bytes] = b"", status_code: int = 500 ) -> None: super().__init__(message, status_code=status_code) self.stdout = stdout or b"" self.stderr = stderr or b"" self.return_code = return_code class SCProcess(asyncio.subprocess.Process): def initialize(self, program_name: str, std_out_cb: OutputCallback, std_err_cb: OutputCallback, log_stderr: bool ) -> None: self.program_name = program_name self.std_out_cb = std_out_cb self.std_err_cb = std_err_cb self.log_stderr = log_stderr self.cancel_requested = False async def _read_stream_with_cb(self, fd: int) -> bytes: transport: BaseTransport = \ self._transport.get_pipe_transport(fd) # type: ignore if fd == 2: stream = self.stderr cb = self.std_err_cb else: assert fd == 1 stream = self.stdout cb = self.std_out_cb assert stream is not None while not stream.at_eof(): output = await stream.readline() if not output: break if fd == 2 and self.log_stderr: logging.info(f"{self.program_name}: {output.decode()}") output = output.rstrip(b'\n') if output and cb is not None: cb(output) transport.close() return output async def cancel(self, sig_idx: int = 1) -> None: if self.cancel_requested: return self.cancel_requested = True exit_success = False sig_idx = min(2, max(0, sig_idx)) sigs = [signal.SIGINT, signal.SIGTERM, signal.SIGKILL][sig_idx:] for sig in sigs: self.send_signal(sig) try: ret = self.wait() await asyncio.wait_for(ret, timeout=2.) except asyncio.TimeoutError: continue logging.debug(f"Command '{self.program_name}' exited with " f"signal: {sig.name}") exit_success = True break if not exit_success: logging.info(f"WARNING: {self.program_name} did not cleanly exit") if self.stdout is not None: self.stdout.feed_eof() if self.stderr is not None: self.stderr.feed_eof() async def communicate_with_cb(self, input: Optional[bytes] = None ) -> None: if input is not None: stdin: Coroutine = self._feed_stdin(input) # type: ignore else: stdin = self._noop() # type: ignore if self.stdout is not None and self.std_out_cb is not None: stdout: Coroutine = self._read_stream_with_cb(1) else: stdout = self._noop() # type: ignore has_err_output = self.std_err_cb is not None or self.log_stderr if self.stderr is not None and has_err_output: stderr: Coroutine = self._read_stream_with_cb(2) else: stderr = self._noop() # type: ignore stdin, stdout, stderr = await asyncio.tasks.gather( stdin, stdout, stderr, loop=self._loop) # type: ignore await self.wait() class ShellCommand: IDX_SIGINT = 0 IDX_SIGTERM = 1 IDX_SIGKILL = 2 def __init__(self, cmd: str, std_out_callback: OutputCallback, std_err_callback: OutputCallback, env: Optional[Dict[str, str]] = None, log_stderr: bool = False, cwd: Optional[str] = None ) -> None: self.name = cmd self.std_out_cb = std_out_callback self.std_err_cb = std_err_callback cmd = os.path.expanduser(cmd) self.command = shlex.split(cmd) self.log_stderr = log_stderr self.env = env self.cwd = cwd self.proc: Optional[SCProcess] = None self.cancelled = False self.return_code: Optional[int] = None async def cancel(self, sig_idx: int = 1) -> None: self.cancelled = True if self.proc is not None: await self.proc.cancel(sig_idx) def get_return_code(self) -> Optional[int]: return self.return_code def _reset_command_data(self) -> None: self.return_code = self.proc = None self.cancelled = False async def run(self, timeout: float = 2., verbose: bool = True, log_complete: bool = True, sig_idx: int = 1 ) -> bool: self._reset_command_data() if not timeout: # Never timeout timeout = 9999999999999999. if self.std_out_cb is None and self.std_err_cb is None and \ not self.log_stderr: # No callbacks set so output cannot be verbose verbose = False if not await self._create_subprocess(): return False assert self.proc is not None try: if verbose: ret: Coroutine = self.proc.communicate_with_cb() else: ret = self.proc.wait() await asyncio.wait_for(ret, timeout=timeout) except asyncio.TimeoutError: complete = False await self.proc.cancel(sig_idx) else: complete = not self.cancelled return self._check_proc_success(complete, log_complete) async def run_with_response(self, timeout: float = 2., retries: int = 1, log_complete: bool = True, sig_idx: int = 1 ) -> str: self._reset_command_data() retries = max(1, retries) while retries > 0: stdout = stderr = b"" if await self._create_subprocess(): assert self.proc is not None try: ret = self.proc.communicate() stdout, stderr = await asyncio.wait_for( ret, timeout=timeout) except asyncio.TimeoutError: complete = False await self.proc.cancel(sig_idx) else: complete = not self.cancelled if self.log_stderr and stderr: logging.info(f"{self.command[0]}: {stderr.decode()}") if self._check_proc_success(complete, log_complete): return stdout.decode().rstrip("\n") elif stdout: logging.debug( f"Shell command '{self.name}' output:" f"\n{stdout.decode()}") retries -= 1 await gen.sleep(.5) raise ShellCommandError( f"Error running shell command: '{self.command}'", self.return_code, stdout, stderr) async def _create_subprocess(self) -> bool: loop = asyncio.get_event_loop() def protocol_factory(): return asyncio.subprocess.SubprocessStreamProtocol( limit=2**20, loop=loop) try: if self.std_err_cb is not None or self.log_stderr: errpipe = asyncio.subprocess.PIPE else: errpipe = asyncio.subprocess.STDOUT transport, protocol = await loop.subprocess_exec( protocol_factory, *self.command, stdout=asyncio.subprocess.PIPE, stderr=errpipe, env=self.env, cwd=self.cwd) self.proc = SCProcess(transport, protocol, loop) self.proc.initialize(self.command[0], self.std_out_cb, self.std_err_cb, self.log_stderr) except Exception: logging.exception( f"shell_command: Command ({self.name}) failed") return False return True def _check_proc_success(self, complete: bool, log_complete: bool ) -> bool: assert self.proc is not None self.return_code = self.proc.returncode success = self.return_code == 0 and complete if success: msg = f"Command ({self.name}) successfully finished" elif self.cancelled: msg = f"Command ({self.name}) cancelled" elif not complete: msg = f"Command ({self.name}) timed out" else: msg = f"Command ({self.name}) exited with return code" \ f" {self.return_code}" if log_complete: logging.info(msg) return success class ShellCommandFactory: error = ShellCommandError def build_shell_command(self, cmd: str, callback: OutputCallback = None, std_err_callback: OutputCallback = None, env: Optional[Dict[str, str]] = None, log_stderr: bool = False, cwd: Optional[str] = None ) -> ShellCommand: return ShellCommand(cmd, callback, std_err_callback, env, log_stderr, cwd) def load_component(config: ConfigHelper) -> ShellCommandFactory: return ShellCommandFactory()