From 544e6fe9de27c89e4248ef08a5c2c8bc38e63aa1 Mon Sep 17 00:00:00 2001 From: Antonio Esposito Date: Thu, 12 Apr 2018 12:18:10 +0200 Subject: [PATCH] Fix #8 Add basic stdin support to Subprocess and Paramiko --- README.rst | 1 + doc/source/ExecResult.rst | 9 ++- doc/source/SSHClient.rst | 5 +- doc/source/Subprocess.rst | 10 ++-- exec_helpers/_ssh_client_base.py | 18 ++++-- exec_helpers/exec_result.py | 16 +++++- exec_helpers/subprocess_runner.py | 13 ++++- test/test_ssh_client.py | 94 +++++++++++++++++++++++++++++++ test/test_subprocess_runner.py | 94 +++++++++++++++++++++++++++++++ 9 files changed, 247 insertions(+), 13 deletions(-) diff --git a/README.rst b/README.rst index 3651ff0..97bdb6c 100644 --- a/README.rst +++ b/README.rst @@ -181,6 +181,7 @@ Execution result object has a set of useful properties: * `cmd` - Command * `exit_code` - Command return code. If possible to decode using enumerators for Linux -> it used. +* `stdin` -> `str`. Text representation of stdin. * `stdout` -> `typing.Tuple[bytes]`. Raw stdout output. * `stderr` -> `typing.Tuple[bytes]`. Raw stderr output. * `stdout_bin` -> `bytearray`. Binary stdout output. diff --git a/doc/source/ExecResult.rst b/doc/source/ExecResult.rst index 6330ae3..05ced46 100644 --- a/doc/source/ExecResult.rst +++ b/doc/source/ExecResult.rst @@ -10,10 +10,12 @@ API: ExecResult Command execution result. - .. py:method:: __init__(cmd, stdout=None, stderr=None, exit_code=ExitCodes.EX_INVALID) + .. py:method:: __init__(cmd, stdin=None, stdout=None, stderr=None, exit_code=ExitCodes.EX_INVALID) :param cmd: command :type cmd: ``str`` + :param stdin: STDIN + :type stdin: ``typing.Optional[str]`` :param stdout: binary STDOUT :type stdout: ``typing.Optional[typing.Iterable[bytes]]`` :param stderr: binary STDERR @@ -36,6 +38,11 @@ API: ExecResult ``str`` Command + .. py:attribute:: stdin + + ``str`` + Stdin input as string. + .. py:attribute:: stdout ``typing.Tuple[bytes]`` diff --git a/doc/source/SSHClient.rst b/doc/source/SSHClient.rst index 73d3bee..28f030d 100644 --- a/doc/source/SSHClient.rst +++ b/doc/source/SSHClient.rst @@ -101,7 +101,7 @@ API: SSHClient and SSHAuth. :param enforce: Enforce sudo enabled or disabled. By default: None :type enforce: ``typing.Optional[bool]`` - .. py:method:: execute_async(command, get_pty=False, open_stdout=True, open_stderr=True, **kwargs) + .. py:method:: execute_async(command, get_pty=False, open_stdout=True, open_stderr=True, stdin=None, **kwargs) Execute command in async mode and return channel with IO objects. @@ -109,6 +109,8 @@ API: SSHClient and SSHAuth. :type command: ``str`` :param get_pty: open PTY on remote machine :type get_pty: ``bool`` + :param stdin: pass STDIN text to the process + :type stdin: ``typing.Union[six.text_type, six.binary_type, None]`` :param open_stdout: open STDOUT stream for read :type open_stdout: bool :param open_stderr: open STDERR stream for read @@ -116,6 +118,7 @@ API: SSHClient and SSHAuth. :rtype: ``typing.Tuple[paramiko.Channel, paramiko.ChannelFile, paramiko.ChannelFile, paramiko.ChannelFile]`` .. versionchanged:: 1.2.0 open_stdout and open_stderr flags + .. versionchanged:: 1.2.0 stdin data .. py:method:: execute(command, verbose=False, timeout=1*60*60, **kwargs) diff --git a/doc/source/Subprocess.rst b/doc/source/Subprocess.rst index f883e06..9aed138 100644 --- a/doc/source/Subprocess.rst +++ b/doc/source/Subprocess.rst @@ -45,6 +45,8 @@ API: Subprocess :param command: Command for execution :type command: ``str`` + :param stdin: STDIN passed to execution + :type stdin: ``typing.Union[six.text_type, six.binary_type, None]`` :param verbose: Produce log.info records for command call and output :type verbose: ``bool`` :param timeout: Timeout for command execution. @@ -52,11 +54,11 @@ API: Subprocess :rtype: ExecResult :raises ExecHelperTimeoutError: Timeout exceeded + .. note:: stdin channel is closed after the input processing .. versionchanged:: 1.1.0 make method - .. versionchanged:: 1.2.0 - - open_stdout and open_stderr flags - default timeout 1 hour + .. versionchanged:: 1.2.0 open_stdout and open_stderr flags + .. versionchanged:: 1.2.0 default timeout 1 hour + .. versionchanged:: 1.2.0 stdin data .. py:method:: check_call(command, verbose=False, timeout=1*60*60, error_info=None, expected=None, raise_on_err=True, **kwargs) diff --git a/exec_helpers/_ssh_client_base.py b/exec_helpers/_ssh_client_base.py index 8015cba..074d4ae 100644 --- a/exec_helpers/_ssh_client_base.py +++ b/exec_helpers/_ssh_client_base.py @@ -489,6 +489,7 @@ def execute_async( self, command, # type: str get_pty=False, # type: bool + stdin=None, # type: typing.Union[six.text_type, six.binary_type, None] open_stdout=True, # type: bool open_stderr=True, # type: bool **kwargs @@ -499,6 +500,8 @@ def execute_async( :type command: str :param get_pty: open PTY on remote machine :type get_pty: bool + :param stdin: pass STDIN text to the process + :type stdin: typing.Union[six.text_type, six.binary_type, None] :param open_stdout: open STDOUT stream for read :type open_stdout: bool :param open_stderr: open STDERR stream for read @@ -511,6 +514,7 @@ def execute_async( ] .. versionchanged:: 1.2.0 open_stdout and open_stderr flags + .. versionchanged:: 1.2.0 stdin data """ cmd_for_log = self._mask_command( cmd=command, @@ -532,7 +536,7 @@ def execute_async( width_pixels=0, height_pixels=0 ) - stdin = chan.makefile('wb') + _stdin = chan.makefile('wb') stdout = chan.makefile('rb') if open_stdout else None stderr = chan.makefile_stderr('rb') if open_stderr else None cmd = "{command}\n".format(command=command) @@ -545,11 +549,17 @@ def execute_async( ) chan.exec_command(cmd) # nosec # Sanitize on caller side if stdout.channel.closed is False: - self.auth.enter_password(stdin) - stdin.flush() + self.auth.enter_password(_stdin) + _stdin.flush() else: chan.exec_command(cmd) # nosec # Sanitize on caller side - return chan, stdin, stderr, stdout + if stdin is not None: + if not isinstance(stdin, six.binary_type): + stdin = stdin.encode(encoding='utf-8') + _stdin.write('{}\n'.format(stdin)) + _stdin.flush() + + return chan, _stdin, stderr, stdout def __exec_command( self, diff --git a/exec_helpers/exec_result.py b/exec_helpers/exec_result.py index a82c8fa..924497d 100644 --- a/exec_helpers/exec_result.py +++ b/exec_helpers/exec_result.py @@ -42,7 +42,7 @@ class ExecResult(object): """Execution result.""" __slots__ = [ - '__cmd', '__stdout', '__stderr', '__exit_code', + '__cmd', '__stdin', '__stdout', '__stderr', '__exit_code', '__timestamp', '__stdout_str', '__stderr_str', '__stdout_brief', '__stderr_brief', '__lock' @@ -51,6 +51,7 @@ class ExecResult(object): def __init__( self, cmd, # type: str + stdin=None, # type: typing.Union[six.text_type, six.binary_type, None] stdout=None, # type: typing.Optional[typing.Iterable[bytes]] stderr=None, # type: typing.Optional[typing.Iterable[bytes]] exit_code=proc_enums.ExitCodes.EX_INVALID # type: _type_exit_codes @@ -59,6 +60,8 @@ def __init__( :param cmd: command :type cmd: str + :param stdin: string STDIN + :type stdin: typing.Union[six.text_type, six.binary_type, None] :param stdout: binary STDOUT :type stdout: typing.Optional[typing.Iterable[bytes]] :param stderr: binary STDERR @@ -69,6 +72,9 @@ def __init__( self.__lock = threading.RLock() self.__cmd = cmd + if stdin is not None and not isinstance(stdin, six.text_type): + stdin = self._get_str_from_bin(stdin) + self.__stdin = stdin self.__stdout = tuple(stdout) if stdout is not None else () self.__stderr = tuple(stderr) if stderr is not None else () @@ -141,6 +147,14 @@ def cmd(self): # type: () -> str """ return self.__cmd + @property + def stdin(self): # type: () -> str + """Stdin input as string. + + :rtype: str + """ + return self.__stdin + @property def stdout(self): # type: () -> typing.Tuple[bytes] """Stdout output as list of binaries. diff --git a/exec_helpers/subprocess_runner.py b/exec_helpers/subprocess_runner.py index 4730fab..49492f2 100644 --- a/exec_helpers/subprocess_runner.py +++ b/exec_helpers/subprocess_runner.py @@ -169,6 +169,7 @@ def __exec_command( timeout=constants.DEFAULT_TIMEOUT, # type: typing.Optional[int] verbose=False, # type: bool log_mask_re=None, # type: typing.Optional[str] + stdin=None, # type: typing.Union[six.text_type, six.binary_type, None] open_stdout=True, # type: bool open_stderr=True, # type: bool ): @@ -183,6 +184,7 @@ def __exec_command( :param log_mask_re: regex lookup rule to mask command for logger. all MATCHED groups will be replaced by '<*masked*>' :type log_mask_re: typing.Optional[str] + :type stdin: typing.Union[six.text_type, six.binary_type, None] :param open_stdout: open STDOUT stream for read :type open_stdout: bool :param open_stderr: open STDERR stream for read @@ -277,12 +279,19 @@ def poll_pipes( # Run self.__process = subprocess.Popen( args=[command], - stdin=subprocess.PIPE, stdout=subprocess.PIPE if open_stdout else devnull, stderr=subprocess.PIPE if open_stderr else devnull, - shell=True, cwd=cwd, env=env, + stdin=subprocess.PIPE, + shell=True, + cwd=cwd, + env=env, universal_newlines=False, ) + if stdin is not None: + if not isinstance(stdin, six.binary_type): + stdin = stdin.encode(encoding='utf-8') + self.__process.stdin.write(stdin) + self.__process.stdin.close() # Poll output diff --git a/test/test_ssh_client.py b/test/test_ssh_client.py index 010b83d..1b2372a 100644 --- a/test/test_ssh_client.py +++ b/test/test_ssh_client.py @@ -61,6 +61,7 @@ def __iter__(self): encoded_cmd = base64.b64encode( "{}\n".format(command).encode('utf-8') ).decode('utf-8') +print_stdin = 'read line; echo "$line"' @mock.patch('exec_helpers._ssh_client_base.logger', autospec=True) @@ -1046,6 +1047,99 @@ def test_check_stderr(self, check_call, client, policy, logger): command, verbose, timeout=None, error_info=None, raise_on_err=raise_on_err) + @mock.patch('exec_helpers.ssh_client.SSHClient.check_call') + def test_check_stdin_str(self, check_call, client, policy, logger): + stdin = u'this is a line' + + return_value = exec_result.ExecResult( + cmd=print_stdin, + stdin=stdin, + stdout=[stdin], + stderr=[], + exit_code=0 + ) + check_call.return_value = return_value + + verbose = False + raise_on_err = True + + # noinspection PyTypeChecker + result = self.get_ssh().check_call( + command=print_stdin, + stdin=stdin, + verbose=verbose, + timeout=None, + raise_on_err=raise_on_err) + check_call.assert_called_once_with( + command=print_stdin, + stdin=stdin, + verbose=verbose, + timeout=None, + raise_on_err=raise_on_err) + self.assertEqual(result, return_value) + + @mock.patch('exec_helpers.ssh_client.SSHClient.check_call') + def test_check_stdin_bytes(self, check_call, client, policy, logger): + stdin = b'this is a line' + + return_value = exec_result.ExecResult( + cmd=print_stdin, + stdin=stdin, + stdout=[stdin], + stderr=[], + exit_code=0 + ) + check_call.return_value = return_value + + verbose = False + raise_on_err = True + + # noinspection PyTypeChecker + result = self.get_ssh().check_call( + command=print_stdin, + stdin=stdin, + verbose=verbose, + timeout=None, + raise_on_err=raise_on_err) + check_call.assert_called_once_with( + command=print_stdin, + stdin=stdin, + verbose=verbose, + timeout=None, + raise_on_err=raise_on_err) + self.assertEqual(result, return_value) + + @mock.patch('exec_helpers.ssh_client.SSHClient.check_call') + def test_check_stdin_bytearray(self, check_call, client, policy, logger): + stdin = bytearray(b'this is a line') + + return_value = exec_result.ExecResult( + cmd=print_stdin, + stdin=stdin, + stdout=[stdin], + stderr=[], + exit_code=0 + ) + check_call.return_value = return_value + + verbose = False + raise_on_err = True + + # noinspection PyTypeChecker + result = self.get_ssh().check_call( + command=print_stdin, + stdin=stdin, + verbose=verbose, + timeout=None, + raise_on_err=raise_on_err) + check_call.assert_called_once_with( + command=print_stdin, + stdin=stdin, + verbose=verbose, + timeout=None, + raise_on_err=raise_on_err) + self.assertEqual(result, return_value) + @mock.patch('exec_helpers._ssh_client_base.logger', autospec=True) @mock.patch( diff --git a/test/test_subprocess_runner.py b/test/test_subprocess_runner.py index 11c3d85..28f0b6d 100644 --- a/test/test_subprocess_runner.py +++ b/test/test_subprocess_runner.py @@ -33,6 +33,7 @@ command_log = u"Executing command:\n{!s}\n".format(command.rstrip()) stdout_list = [b' \n', b'2\n', b'3\n', b' \n'] stderr_list = [b' \n', b'0\n', b'1\n', b' \n'] +print_stdin = 'read line; echo "$line"' class FakeFileStream(object): @@ -577,3 +578,96 @@ def test_check_stderr(self, check_call, logger): check_call.assert_called_once_with( command, verbose, timeout=None, error_info=None, raise_on_err=raise_on_err) + + @mock.patch('exec_helpers.subprocess_runner.Subprocess.check_call') + def test_check_stdin_str(self, check_call, logger): + stdin = u'this is a line' + + expected_result = exec_helpers.ExecResult( + cmd=print_stdin, + stdin=stdin, + stdout=[stdin], + stderr=[b''], + exit_code=0, + ) + check_call.return_value = expected_result + + verbose = False + + runner = exec_helpers.Subprocess() + + # noinspection PyTypeChecker + result = runner.check_call( + command=print_stdin, + verbose=verbose, + timeout=None, + stdin=stdin) + check_call.assert_called_once_with( + command=print_stdin, + verbose=verbose, + timeout=None, + stdin=stdin) + self.assertEqual(result, expected_result) + assert result == expected_result + + @mock.patch('exec_helpers.subprocess_runner.Subprocess.check_call') + def test_check_stdin_bytes(self, check_call, logger): + stdin = b'this is a line' + + expected_result = exec_helpers.ExecResult( + cmd=print_stdin, + stdin=stdin, + stdout=[stdin], + stderr=[b''], + exit_code=0, + ) + check_call.return_value = expected_result + + verbose = False + + runner = exec_helpers.Subprocess() + + # noinspection PyTypeChecker + result = runner.check_call( + command=print_stdin, + verbose=verbose, + timeout=None, + stdin=stdin) + check_call.assert_called_once_with( + command=print_stdin, + verbose=verbose, + timeout=None, + stdin=stdin) + self.assertEqual(result, expected_result) + assert result == expected_result + + @mock.patch('exec_helpers.subprocess_runner.Subprocess.check_call') + def test_check_stdin_bytearray(self, check_call, logger): + stdin = bytearray(b'this is a line') + + expected_result = exec_helpers.ExecResult( + cmd=print_stdin, + stdin=stdin, + stdout=[stdin], + stderr=[b''], + exit_code=0, + ) + check_call.return_value = expected_result + + verbose = False + + runner = exec_helpers.Subprocess() + + # noinspection PyTypeChecker + result = runner.check_call( + command=print_stdin, + verbose=verbose, + timeout=None, + stdin=stdin) + check_call.assert_called_once_with( + command=print_stdin, + verbose=verbose, + timeout=None, + stdin=stdin) + self.assertEqual(result, expected_result) + assert result == expected_result