Support chroot with property, context and argument

Signed-off-by: Aleksei Stepanov <penguinolog@gmail.com>

Author: penguinolog

doc/source/SSHClient.rst

@@ -10,7 +10,7 @@ API: SSHClient and SSHAuth.
 
     SSHClient helper.
 
-    .. py:method:: __init__(host, port=22, username=None, password=None, private_keys=None, auth=None, )
+    .. py:method:: __init__(host, port=22, username=None, password=None, private_keys=None, auth=None, *, chroot_path=None)
 
         :param host: remote hostname
         :type host: ``str``
@@ -26,6 +26,8 @@ API: SSHClient and SSHAuth.
         :type auth: typing.Optional[SSHAuth]
         :param verbose: show additional error/warning messages
         :type verbose: bool
+        :param chroot_path: chroot path (use chroot if set)
+        :type chroot_path: typing.Optional[str]
 
     .. note:: auth has priority over username/password/private_keys
 
@@ -66,6 +68,11 @@ API: SSHClient and SSHAuth.
         ``bool``
         Paramiko status: ready to use|reconnect required
 
+    .. py:attribute:: chroot_path
+
+        ``typing.Optional[str]``
+        Path for chroot if set.
+
     .. py:attribute:: sudo_mode
 
         ``bool``
@@ -102,6 +109,18 @@ API: SSHClient and SSHAuth.
         .. versionchanged:: 1.1.0 release lock on exit
         .. versionchanged:: 1.2.1 disconnect enforced on close only not in keepalive mode
 
+    .. py:method:: chroot(path)
+
+        Context manager for changing chroot rules.
+
+        :param path: chroot path or none for working without chroot.
+        :type path: typing.Optional[str]
+        :return: context manager with selected chroot state inside
+        :rtype: typing.ContextManager
+
+        .. Note:: Enter and exit main context manager is produced as well.
+        .. versionadded:: 4.1.0
+
     .. py:method:: sudo(enforce=None)
 
         Context manager getter for sudo operation
@@ -121,7 +140,7 @@ API: SSHClient and SSHAuth.
         .. Note:: Enter and exit ssh context manager is produced as well.
         .. versionadded:: 1.2.1
 
-    .. py:method:: execute_async(command, stdin=None, open_stdout=True, open_stderr=True, verbose=False, log_mask_re=None, *, get_pty=False, width=80, height=24, **kwargs)
+    .. py:method:: execute_async(command, stdin=None, open_stdout=True, open_stderr=True, verbose=False, log_mask_re=None, *, chroot_path=None, get_pty=False, width=80, height=24, **kwargs)
 
         Execute command in async mode and return channel with IO objects.
 
@@ -137,7 +156,9 @@ API: SSHClient and SSHAuth.
         :type verbose: bool
         :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 log_mask_re: ``typing.Optional[str]``
+        :param chroot_path: chroot path override
+        :type chroot_path: ``typing.Optional[str]``
         :param get_pty: Get PTY for connection
         :type get_pty: bool
         :param width: PTY width

doc/source/Subprocess.rst

@@ -8,18 +8,21 @@ API: Subprocess
 
 .. py:class:: Subprocess()
 
-    .. py:method:: __init__(logger, log_mask_re=None, *, logger=logging.getLogger("exec_helpers.subprocess_runner"))
+    .. py:method:: __init__(logger, log_mask_re=None, *, logger=logging.getLogger("exec_helpers.subprocess_runner"), chroot_path=None)
 
         ExecHelper global API.
 
         :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]
         :param logger: logger instance to use
         :type logger: logging.Logger
+        :param chroot_path: chroot path (use chroot if set)
+        :type chroot_path: typing.Optional[str]
 
         .. versionchanged:: 1.2.0 log_mask_re regex rule for masking cmd
         .. versionchanged:: 3.1.0 Not singleton anymore. Only lock is shared between all instances.
         .. versionchanged:: 3.2.0 Logger can be enforced.
+        .. versionchanged:: 4.1.0 support chroot
 
     .. py:attribute:: log_mask_re
 
@@ -43,7 +46,24 @@ API: Subprocess
 
         .. versionchanged:: 1.1.0 release lock on exit
 
-    .. py:method:: execute_async(command, stdin=None, open_stdout=True, open_stderr=True, verbose=False, log_mask_re=None, *, cwd=None, env=None, **kwargs)
+    .. py:attribute:: chroot_path
+
+        ``typing.Optional[str]``
+        Path for chroot if set.
+
+    .. py:method:: chroot(path)
+
+        Context manager for changing chroot rules.
+
+        :param path: chroot path or none for working without chroot.
+        :type path: typing.Optional[str]
+        :return: context manager with selected chroot state inside
+        :rtype: typing.ContextManager
+
+        .. Note:: Enter and exit main context manager is produced as well.
+        .. versionadded:: 4.1.0
+
+    .. py:method:: execute_async(command, stdin=None, open_stdout=True, open_stderr=True, verbose=False, log_mask_re=None, *, chroot_path=None, cwd=None, env=None, **kwargs)
 
         Execute command in async mode and return Popen with IO objects.
 
@@ -60,6 +80,8 @@ API: Subprocess
         :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]``
+        :param chroot_path: chroot path override
+        :type chroot_path: ``typing.Optional[str]``
         :param cwd: Sets the current directory before the child is executed.
         :type cwd: typing.Optional[typing.Union[str, bytes]]
         :param env: Defines the environment variables for the new process.

exec_helpers/_ssh_client_base.py

@@ -128,6 +128,8 @@ def __call__(  # type: ignore
         private_keys: typing.Optional[typing.Iterable[paramiko.RSAKey]] = None,
         auth: typing.Optional[ssh_auth.SSHAuth] = None,
         verbose: bool = True,
+        *,
+        chroot_path: typing.Optional[str] = None,
     ) -> "SSHClientBase":
         """Main memorize method: check for cached instance and return it. API follows target __init__.
 
@@ -145,10 +147,12 @@ def __call__(  # type: ignore
         :type auth: typing.Optional[ssh_auth.SSHAuth]
         :param verbose: show additional error/warning messages
         :type verbose: bool
+        :param chroot_path: chroot path (use chroot if set)
+        :type chroot_path: typing.Optional[str]
         :return: SSH client instance
         :rtype: SSHClientBase
         """
-        if (host, port) in cls.__cache:
+        if (host, port) in cls.__cache and not chroot_path:  # chrooted connections are not memorized
             key = host, port
             if auth is None:
                 auth = ssh_auth.SSHAuth(username=username, password=password, keys=private_keys)
@@ -176,8 +180,10 @@ def __call__(  # type: ignore
             private_keys=private_keys,
             auth=auth,
             verbose=verbose,
+            chroot_path=chroot_path,
         )
-        cls.__cache[(ssh.hostname, ssh.port)] = ssh
+        if not chroot_path:
+            cls.__cache[(ssh.hostname, ssh.port)] = ssh
         return ssh
 
     @classmethod
@@ -188,7 +194,6 @@ def clear_cache(mcs: typing.Type["_MemorizedSSH"]) -> None:
         """
         n_count = 3
         # PY3: cache, ssh, temporary
-        # PY4: cache, values mapping, ssh, temporary
         for ssh in mcs.__cache.values():
             if CPYTHON and sys.getrefcount(ssh) == n_count:  # pragma: no cover
                 ssh.logger.debug("Closing as unused")
@@ -203,63 +208,66 @@ def close_connections(mcs: typing.Type["_MemorizedSSH"]) -> None:
                 ssh.close()  # type: ignore
 
 
-class SSHClientBase(api.ExecHelper, metaclass=_MemorizedSSH):
-    """SSH Client helper."""
+class _SudoContext:
+    """Context manager for call commands with sudo."""
 
-    __slots__ = ("__hostname", "__port", "__auth", "__ssh", "__sftp", "__sudo_mode", "__keepalive_mode", "__verbose")
+    __slots__ = ("__ssh", "__sudo_status", "__enforce")
 
-    class __get_sudo:
-        """Context manager for call commands with sudo."""
+    def __init__(self, ssh: "SSHClientBase", enforce: typing.Optional[bool] = None) -> None:
+        """Context manager for call commands with sudo.
 
-        __slots__ = ("__ssh", "__sudo_status", "__enforce")
+        :param ssh: connection instance
+        :type ssh: SSHClientBase
+        :param enforce: sudo mode for context manager
+        :type enforce: typing.Optional[bool]
+        """
+        self.__ssh: "SSHClientBase" = ssh
+        self.__sudo_status: bool = ssh.sudo_mode
+        self.__enforce: typing.Optional[bool] = enforce
 
-        def __init__(self, ssh: "SSHClientBase", enforce: typing.Optional[bool] = None) -> None:
-            """Context manager for call commands with sudo.
+    def __enter__(self) -> None:
+        self.__sudo_status = self.__ssh.sudo_mode
+        if self.__enforce is not None:
+            self.__ssh.sudo_mode = self.__enforce
 
-            :param ssh: connection instance
-            :type ssh: SSHClientBase
-            :param enforce: sudo mode for context manager
-            :type enforce: typing.Optional[bool]
-            """
-            self.__ssh: "SSHClientBase" = ssh
-            self.__sudo_status: bool = ssh.sudo_mode
-            self.__enforce: typing.Optional[bool] = enforce
+    def __exit__(self, exc_type: typing.Any, exc_val: typing.Any, exc_tb: typing.Any) -> None:
+        self.__ssh.sudo_mode = self.__sudo_status
 
-        def __enter__(self) -> None:
-            self.__sudo_status = self.__ssh.sudo_mode
-            if self.__enforce is not None:
-                self.__ssh.sudo_mode = self.__enforce
 
-        def __exit__(self, exc_type: typing.Any, exc_val: typing.Any, exc_tb: typing.Any) -> None:
-            self.__ssh.sudo_mode = self.__sudo_status
+class _KeepAliveContext:
+    """Context manager for keepalive management."""
 
-    class __get_keepalive:
-        """Context manager for keepalive management."""
+    __slots__ = ("__ssh", "__keepalive_status", "__enforce")
 
-        __slots__ = ("__ssh", "__keepalive_status", "__enforce")
+    def __init__(self, ssh: "SSHClientBase", enforce: bool = True) -> None:
+        """Context manager for keepalive management.
 
-        def __init__(self, ssh: "SSHClientBase", enforce: bool = True) -> None:
-            """Context manager for keepalive management.
+        :param ssh: connection instance
+        :type ssh: SSHClientBase
+        :param enforce: keepalive mode for context manager
+        :type enforce: bool
+        :param enforce: Keep connection alive after context manager exit
+        """
+        self.__ssh: "SSHClientBase" = ssh
+        self.__keepalive_status: bool = ssh.keepalive_mode
+        self.__enforce: typing.Optional[bool] = enforce
 
-            :param ssh: connection instance
-            :type ssh: SSHClientBase
-            :param enforce: keepalive mode for context manager
-            :type enforce: bool
-            :param enforce: Keep connection alive after context manager exit
-            """
-            self.__ssh: "SSHClientBase" = ssh
-            self.__keepalive_status: bool = ssh.keepalive_mode
-            self.__enforce: typing.Optional[bool] = enforce
+    def __enter__(self) -> None:
+        self.__ssh.__enter__()
+        self.__keepalive_status = self.__ssh.keepalive_mode
+        if self.__enforce is not None:
+            self.__ssh.keepalive_mode = self.__enforce
 
-        def __enter__(self) -> None:
-            self.__keepalive_status = self.__ssh.keepalive_mode
-            if self.__enforce is not None:
-                self.__ssh.keepalive_mode = self.__enforce
-            self.__ssh.__enter__()
+    def __exit__(self, exc_type: typing.Any, exc_val: typing.Any, exc_tb: typing.Any) -> None:
+        # Exit before releasing!
+        self.__ssh.__exit__(exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb)  # type: ignore
+        self.__ssh.keepalive_mode = self.__keepalive_status
+
+
+class SSHClientBase(api.ExecHelper, metaclass=_MemorizedSSH):
+    """SSH Client helper."""
 
-        def __exit__(self, exc_type: typing.Any, exc_val: typing.Any, exc_tb: typing.Any) -> None:
-            self.__ssh.__exit__(exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb)  # type: ignore
-            self.__ssh.keepalive_mode = self.__keepalive_status
+    __slots__ = ("__hostname", "__port", "__auth", "__ssh", "__sftp", "__sudo_mode", "__keepalive_mode", "__verbose")
 
     def __hash__(self) -> int:
         """Hash for usage as dict keys."""
@@ -274,6 +282,8 @@ def __init__(
         private_keys: typing.Optional[typing.Iterable[paramiko.RSAKey]] = None,
         auth: typing.Optional[ssh_auth.SSHAuth] = None,
         verbose: bool = True,
+        *,
+        chroot_path: typing.Optional[str] = None,
     ) -> None:
         """Main SSH Client helper.
 
@@ -291,11 +301,13 @@ def __init__(
         :type auth: typing.Optional[ssh_auth.SSHAuth]
         :param verbose: show additional error/warning messages
         :type verbose: bool
+        :param chroot_path: chroot path (use chroot if set)
+        :type chroot_path: typing.Optional[str]
 
         .. note:: auth has priority over username/password/private_keys
         """
         super(SSHClientBase, self).__init__(
-            logger=logging.getLogger(self.__class__.__name__).getChild(f"{host}:{port}")
+            logger=logging.getLogger(self.__class__.__name__).getChild(f"{host}:{port}"), chroot_path=chroot_path
         )
 
         self.__hostname: str = host
@@ -510,7 +522,7 @@ def sudo(self, enforce: typing.Optional[bool] = None) -> "typing.ContextManager[
         :return: context manager with selected sudo state inside
         :rtype: typing.ContextManager
         """
-        return self.__get_sudo(ssh=self, enforce=enforce)
+        return _SudoContext(ssh=self, enforce=enforce)
 
     def keepalive(self, enforce: bool = True) -> "typing.ContextManager[None]":
         """Call contextmanager with keepalive mode change.
@@ -523,7 +535,7 @@ def keepalive(self, enforce: bool = True) -> "typing.ContextManager[None]":
         .. Note:: Enter and exit ssh context manager is produced as well.
         .. versionadded:: 1.2.1
         """
-        return self.__get_keepalive(ssh=self, enforce=enforce)
+        return _KeepAliveContext(ssh=self, enforce=enforce)
 
     # noinspection PyMethodOverriding
     def execute_async(  # pylint: disable=arguments-differ
@@ -535,6 +547,7 @@ def execute_async(  # pylint: disable=arguments-differ
         verbose: bool = False,
         log_mask_re: typing.Optional[str] = None,
         *,
+        chroot_path: typing.Optional[str] = None,
         get_pty: bool = False,
         width: int = 80,
         height: int = 24,
@@ -555,6 +568,8 @@ def execute_async(  # pylint: disable=arguments-differ
         :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]
+        :param chroot_path: chroot path override
+        :type chroot_path: typing.Optional[str]
         :param get_pty: Get PTY for connection
         :type get_pty: bool
         :param width: PTY width
@@ -580,6 +595,7 @@ def execute_async(  # pylint: disable=arguments-differ
         .. versionchanged:: 1.2.0 get_pty moved to `**kwargs`
         .. versionchanged:: 2.1.0 Use typed NamedTuple as result
         .. versionchanged:: 3.2.0 Expose pty options as optional keyword-only arguments
+        .. versionchanged:: 4.1.0 support chroot
         """
         cmd_for_log: str = self._mask_command(cmd=command, log_mask_re=log_mask_re)
 
@@ -597,7 +613,8 @@ def execute_async(  # pylint: disable=arguments-differ
         stdout: paramiko.ChannelFile = chan.makefile("rb")
         stderr: typing.Optional[paramiko.ChannelFile] = chan.makefile_stderr("rb") if open_stderr else None
 
-        cmd = f"{command}\n"
+        cmd = f"{self._prepare_command(cmd=command, chroot_path=chroot_path)}\n"
+
         started = datetime.datetime.utcnow()
         if self.sudo_mode:
             encoded_cmd = base64.b64encode(cmd.encode("utf-8")).decode("utf-8")