Fix #8 Add basic stdin support to Subprocess and Paramiko

Author: espositocloud

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.

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]``

doc/source/SSHClient.rst

@@ -101,21 +101,24 @@ 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.
 
         :param command: Command for execution
         :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
         :type open_stderr: bool
         :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)
 

doc/source/Subprocess.rst

@@ -45,18 +45,20 @@ 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.
         :type timeout: ``typing.Optional[int]``
         :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)
 

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,

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.

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
 

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(