Skip to content

Commit

Permalink
Fix chroot_exe when used with context manager (#156)
Browse files Browse the repository at this point in the history
* Fix chroot_exe when used with context manager

Use None for `chroot_exe` default argument.  This allows us to use
similar logic as `chroot_path`.

* Update exec_helpers/api.py

Co-authored-by: Alexey Stepanov <penguinolog@users.noreply.github.com>

* Update exec_helpers/async_api/api.py

Co-authored-by: Alexey Stepanov <penguinolog@users.noreply.github.com>

---------

Co-authored-by: Geert Kloosterman <geert.kloosterman@brightcomputing.com>
Co-authored-by: Alexey Stepanov <penguinolog@users.noreply.github.com>
  • Loading branch information
3 people committed Jun 18, 2024
1 parent 3129fc4 commit 0b32a78
Show file tree
Hide file tree
Showing 6 changed files with 86 additions and 74 deletions.
5 changes: 3 additions & 2 deletions exec_helpers/_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,18 +78,19 @@ def cmd_to_string(command: str | Iterable[str]) -> str:
return shlex.join(command)


def chroot_command(command: str, chroot_path: str | None = None, chroot_exe: str = "chroot") -> str:
def chroot_command(command: str, chroot_path: str | None = None, chroot_exe: str | None = None) -> str:
"""Prepare command for chroot execution.
:param command: Original command.
:type command: str
:param chroot_path: chroot path.
:type chroot_path: str | None
:param chroot_exe: chroot executable.
:type chroot_exe: str
:type chroot_exe: str | None
:return: Command to be executed with chroot rules if applicable.
:rtype: str
"""
chroot_exe = chroot_exe or "chroot"
if chroot_path and chroot_path != "/":
chroot_dst: str = shlex.quote(chroot_path.strip())
quoted_command = shlex.quote(command)
Expand Down
19 changes: 10 additions & 9 deletions exec_helpers/_ssh_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -982,14 +982,15 @@ def keepalive(self, enforce: KeepAlivePeriodT = 1) -> _KeepAliveContext:
"""
return _KeepAliveContext(ssh=self, enforce=int(enforce))

def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe: str = "chroot") -> str:
def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe: str | None = None) -> str:
"""Prepare command: cover chroot and other cases.
:param cmd: Main command.
:param chroot_path: Path to make chroot for execution.
:param chroot_exe: chroot executable, default "chroot".
:return: The final command includes chroot, if required.
"""
chroot_exe = chroot_exe or "chroot"
if not self.sudo_mode:
return super()._prepare_command(cmd=cmd, chroot_path=chroot_path, chroot_exe=chroot_exe)
quoted_command: str = shlex.quote(cmd)
Expand Down Expand Up @@ -1096,7 +1097,7 @@ def open_execute_context(
open_stdout: bool = True,
open_stderr: bool = True,
chroot_path: str | None = None,
chroot_exe: str = "chroot",
chroot_exe: str | None = None,
get_pty: bool = False,
width: int = 80,
height: int = 24,
Expand All @@ -1116,7 +1117,7 @@ def open_execute_context(
:param chroot_path: chroot path override.
:type chroot_path: str | None
:param chroot_exe: chroot exe override.
:type chroot_exe: str
:type chroot_exe: str | None
:param get_pty: Get PTY for connection.
:type get_pty: bool
:param width: PTY width.
Expand Down Expand Up @@ -1160,7 +1161,7 @@ def execute(
open_stderr: bool = True,
log_stderr: bool = True,
chroot_path: str | None = None,
chroot_exe: str = "chroot",
chroot_exe: str | None = None,
get_pty: bool = False,
width: int = 80,
height: int = 24,
Expand Down Expand Up @@ -1190,7 +1191,7 @@ def execute(
:param chroot_path: chroot path override.
:type chroot_path: str | None
:param chroot_exe: chroot exe override.
:type chroot_exe: str
:type chroot_exe: str | None
:param get_pty: Get PTY for connection.
:type get_pty: bool
:param width: PTY width.
Expand Down Expand Up @@ -1240,7 +1241,7 @@ def __call__(
open_stderr: bool = True,
log_stderr: bool = True,
chroot_path: str | None = None,
chroot_exe: str = "chroot",
chroot_exe: str | None = None,
get_pty: bool = False,
width: int = 80,
height: int = 24,
Expand Down Expand Up @@ -1270,7 +1271,7 @@ def __call__(
:param chroot_path: Chroot path override.
:type chroot_path: str | None
:param chroot_exe: Chroot exe override.
:type chroot_exe: str
:type chroot_exe: str | None
:param get_pty: Get PTY for connection.
:type get_pty: bool
:param width: PTY width.
Expand Down Expand Up @@ -1658,7 +1659,7 @@ def execute_together(
open_stdout: bool = True,
open_stderr: bool = True,
chroot_path: str | None = None,
chroot_exe: str = "chroot",
chroot_exe: str | None = None,
verbose: bool = False,
log_mask_re: LogMaskReT = None,
exception_class: type[exceptions.ParallelCallProcessError] = exceptions.ParallelCallProcessError,
Expand All @@ -1685,7 +1686,7 @@ def execute_together(
:param chroot_path: chroot path override.
:type chroot_path: str | None
:param chroot_exe: chroot exe override.
:type chroot_exe: str
:type chroot_exe: str | None
:param verbose: Produce verbose log record on command call.
:type verbose: bool
:param log_mask_re: Regex lookup rule to mask command for logger.
Expand Down
52 changes: 28 additions & 24 deletions exec_helpers/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,32 +173,32 @@ class _ChRootContext(typing.ContextManager[None]):
:param path: chroot path or None for no chroot.
:type path: str | pathlib.Path | None
:param chroot_exe: chroot executable.
:type chroot_exe: str
:type chroot_exe: str | None
:raises TypeError: Incorrect type of path or chroot_exe variable.
.. versionadded:: 4.1.0
"""

__slots__ = ("_chroot_exe_status", "_chroot_status", "_conn", "_exe", "_path")

def __init__(self, conn: ExecHelper, path: ChRootPathSetT = None, chroot_exe: str = "chroot") -> None:
def __init__(self, conn: ExecHelper, path: ChRootPathSetT = None, chroot_exe: str | None = None) -> None:
"""Context manager for call commands with sudo.
:raises TypeError: Incorrect type of path or chroot_exe variable.
"""
self._conn: ExecHelper = conn
self._chroot_status: str | None = conn._chroot_path
self._chroot_exe_status: str = conn._chroot_exe
self._chroot_exe_status: str | None = conn._chroot_exe
if path is None or isinstance(path, str):
self._path: str | None = path
elif isinstance(path, pathlib.Path):
self._path = path.as_posix() # get absolute path
else:
raise TypeError(f"path={path!r} is not instance of {ChRootPathSetT}")
if isinstance(chroot_exe, str):
self._exe: str = chroot_exe
if chroot_exe is None or isinstance(chroot_exe, str):
self._exe: str | None = chroot_exe
else:
raise TypeError(f"chroot_exe={chroot_exe!r} is not instance of str")
raise TypeError(f"chroot_exe={chroot_exe!r} is not None or instance of str")

def __enter__(self) -> None:
self._conn.__enter__()
Expand Down Expand Up @@ -252,7 +252,7 @@ def __init__(self, log_mask_re: LogMaskReT = None, *, logger: logging.Logger) ->
self.__logger: logging.Logger = logger
self.log_mask_re: LogMaskReT = log_mask_re
self.__chroot_path: str | None = None
self.__chroot_exe: str = "chroot"
self.__chroot_exe: str | None = None
self.__context_count = 0

@property
Expand Down Expand Up @@ -306,43 +306,43 @@ def _chroot_path(self) -> None:
self.__chroot_path = None

@property
def _chroot_exe(self) -> str:
def _chroot_exe(self) -> str | None:
"""Exe for chroot
:rtype: str
:rtype: str | None
.. versionadded:: 8.1.0
"""
return self.__chroot_exe

@_chroot_exe.setter
def _chroot_exe(self, new_state: str) -> None:
def _chroot_exe(self, new_state: str | None) -> None:
"""Executable for chroot if set.
:param new_state: New exe.
:type new_state: str
:type new_state: str | None
:raises TypeError: Not supported exe information.
.. versionadded:: 8.1.0
"""
if isinstance(new_state, str):
if new_state is None or isinstance(new_state, str):
self.__chroot_exe = new_state
else:
raise TypeError(f"chroot_exe is expected to be string, but set {new_state!r}")
raise TypeError(f"chroot_exe is expected to be None or string, but set {new_state!r}")

@_chroot_exe.deleter
def _chroot_exe(self) -> None:
"""Restore chroot executable.
.. versionadded:: 8.1.0
"""
self.__chroot_exe = "chroot"
self.__chroot_exe = None

def chroot(self, path: ChRootPathSetT, chroot_exe: str = "chroot") -> _ChRootContext:
def chroot(self, path: ChRootPathSetT, chroot_exe: str | None = None) -> _ChRootContext:
"""Context manager for changing chroot rules.
:param path: chroot path or none for working without chroot.
:type path: str | pathlib.Path | None
:param chroot_exe: chroot exe.
:type chroot_exe: str
:type chroot_exe: str | None
:return: Context manager with selected chroot state inside.
:rtype: typing.ContextManager
Expand Down Expand Up @@ -394,7 +394,7 @@ def _mask_command(self, cmd: str, log_mask_re: LogMaskReT = None) -> str:

return _helpers.mask_command(cmd.rstrip(), self.log_mask_re, log_mask_re)

def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe: str = "chroot") -> str:
def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe: str | None = None) -> str:
"""Prepare command: cower chroot and other cases.
:param cmd: Main command.
Expand All @@ -404,7 +404,11 @@ def _prepare_command(self, cmd: str, chroot_path: str | None = None, chroot_exe:
:return: Final command, includes chroot, if required.
:rtype: str
"""
return _helpers.chroot_command(cmd, chroot_path=chroot_path or self._chroot_path, chroot_exe=chroot_exe)
return _helpers.chroot_command(
cmd,
chroot_path=chroot_path or self._chroot_path,
chroot_exe=chroot_exe or self._chroot_exe,
)

@abc.abstractmethod
def _exec_command(
Expand Down Expand Up @@ -473,7 +477,7 @@ def open_execute_context(
open_stdout: bool = True,
open_stderr: bool = True,
chroot_path: str | None = None,
chroot_exe: str = "chroot",
chroot_exe: str | None = None,
**kwargs: typing.Any,
) -> ExecuteContext:
"""Get execution context manager.
Expand All @@ -489,7 +493,7 @@ def open_execute_context(
:param chroot_path: chroot path override.
:type chroot_path: str | None
:param chroot_exe: chroot exe override.
:type chroot_exe: str
:type chroot_exe: str | None
:param kwargs: Additional parameters for call.
:type kwargs: typing.Any
Expand All @@ -509,7 +513,7 @@ def execute(
open_stderr: bool = True,
log_stderr: bool = True,
chroot_path: str | None = None,
chroot_exe: str = "chroot",
chroot_exe: str | None = None,
**kwargs: typing.Any,
) -> exec_result.ExecResult:
"""Execute command and wait for return code.
Expand All @@ -536,7 +540,7 @@ def execute(
:param chroot_path: chroot path override.
:type chroot_path: str | None
:param chroot_exe: chroot exe override.
:type chroot_exe: str
:type chroot_exe: str | None
:param kwargs: Additional parameters for call.
:type kwargs: typing.Any
:return: Execution result.
Expand Down Expand Up @@ -594,7 +598,7 @@ def __call__(
open_stderr: bool = True,
log_stderr: bool = True,
chroot_path: str | None = None,
chroot_exe: str = "chroot",
chroot_exe: str | None = None,
**kwargs: typing.Any,
) -> exec_result.ExecResult:
"""Execute command and wait for return code.
Expand All @@ -621,7 +625,7 @@ def __call__(
:param chroot_path: chroot path override.
:type chroot_path: str | None
:param chroot_exe: chroot exe override.
:type chroot_exe: str
:type chroot_exe: str | None
:param kwargs: Additional parameters for call.
:type kwargs: typing.Any
:return: Execution result.
Expand Down
Loading

0 comments on commit 0b32a78

Please sign in to comment.