From 715bc63bb6cb04f2f522e738a753c9e991d7df0e Mon Sep 17 00:00:00 2001 From: Reuben Miller Date: Fri, 1 Nov 2024 09:09:24 +0100 Subject: [PATCH] feat: support setting a default username for command execution --- device_test_core/adapter.py | 13 ++++++++++++- device_test_core/docker/device.py | 6 ++++-- device_test_core/local/device.py | 9 ++++++++- device_test_core/ssh/device.py | 8 +++++++- 4 files changed, 31 insertions(+), 5 deletions(-) diff --git a/device_test_core/adapter.py b/device_test_core/adapter.py index cca4bcb..40dc09e 100644 --- a/device_test_core/adapter.py +++ b/device_test_core/adapter.py @@ -25,6 +25,7 @@ def __init__( should_cleanup: bool = None, use_sudo: bool = True, config: Dict[str, Any] = None, + user: str = "", ): self._name = name self._device_id = ( @@ -36,6 +37,7 @@ def __init__( self._use_sudo = use_sudo self._config = config self._should_cleanup = should_cleanup + self._user = user @property def test_start_time(self) -> datetime: @@ -83,9 +85,17 @@ def use_sudo(self) -> bool: """ return self._use_sudo + def user(self) -> str: + """Shell User + + Returns: + str: user to execute the commands under + """ + return self._user + @abstractmethod def execute_command( - self, cmd: str, log_output: bool = True, shell: bool = True, **kwargs + self, cmd: str, log_output: bool = True, shell: bool = True, user: str = "", **kwargs ) -> CmdOutput: """Execute a command inside the docker container @@ -93,6 +103,7 @@ def execute_command( cmd (str): Command to execute log_output (bool, optional): Log the stdout after the command has executed shell (bool, optional): Execute the command in a shell + user (bool, optional): User that the shell is executed under **kwargs (Any, optional): Additional keyword arguments Returns: diff --git a/device_test_core/docker/device.py b/device_test_core/docker/device.py index e86ac09..c2d3138 100644 --- a/device_test_core/docker/device.py +++ b/device_test_core/docker/device.py @@ -145,7 +145,7 @@ def get_device_stats(self) -> Any: return self.container.stats(stream=False) def execute_command( - self, cmd: str, log_output: bool = True, shell: bool = True, **kwargs + self, cmd: str, log_output: bool = True, shell: bool = True, user: str = "", **kwargs ) -> CmdOutput: """Execute a command inside the docker container @@ -153,6 +153,7 @@ def execute_command( cmd (str): Command to execute log_output (bool, optional): Log the stdout after the command has executed shell (bool, optional): Execute the command in a shell + user (bool, optional): User that the shell is executed under **kwargs (Any, optional): Additional keyword arguments Raises: @@ -162,6 +163,7 @@ def execute_command( CmdOutput: Command output details, e.g. stdout, stderr and return_code """ run_cmd = [] + user = user or self.user() use_sudo = kwargs.pop("sudo", self.use_sudo()) if use_sudo: @@ -176,7 +178,7 @@ def execute_command( else: run_cmd.append(cmd) - exit_code, output = self.container.exec_run(run_cmd, demux=True) + exit_code, output = self.container.exec_run(run_cmd, user=user or self.user(), demux=True) stdout, stderr = output if log_output: log.info( diff --git a/device_test_core/local/device.py b/device_test_core/local/device.py index 4bfc6cd..3ae7eca 100644 --- a/device_test_core/local/device.py +++ b/device_test_core/local/device.py @@ -107,7 +107,7 @@ def get_device_stats(self) -> Any: ) def execute_command( - self, cmd: str, log_output: bool = True, shell: bool = True, **kwargs + self, cmd: str, log_output: bool = True, shell: bool = True, user: str = "", **kwargs ) -> CmdOutput: """Execute a command inside the docker container @@ -115,6 +115,7 @@ def execute_command( cmd (str): Command to execute log_output (bool, optional): Log the stdout after the command has executed shell (bool, optional): Execute the command in a shell + user (bool, optional): User that the shell is executed under **kwargs (Any, optional): Additional keyword arguments Raises: @@ -125,6 +126,12 @@ def execute_command( """ run_cmd = [] + # Note: the local adapter has limited options to change which user is used to run, so use + # sudo to run as a different user + user = user or self.user() + if user: + run_cmd.extend(["sudo", "-E", "-u", user]) + use_sudo = kwargs.pop("sudo", self.use_sudo()) if use_sudo: run_cmd.extend(["sudo", "-E"]) diff --git a/device_test_core/ssh/device.py b/device_test_core/ssh/device.py index d44d36e..991bd5d 100644 --- a/device_test_core/ssh/device.py +++ b/device_test_core/ssh/device.py @@ -182,7 +182,7 @@ def _connect(self): self._client.open() def execute_command( - self, cmd: str, log_output: bool = True, shell: bool = True, **kwargs + self, cmd: str, log_output: bool = True, shell: bool = True, user: str = "", **kwargs ) -> CmdOutput: """Execute a command @@ -190,6 +190,7 @@ def execute_command( cmd (str): Command to execute log_output (bool, optional): Log the stdout after the command has executed shell (bool, optional): Execute the command in a shell + user (bool, optional): User that the shell is executed under **kwargs (Any, optional): Additional keyword arguments Raises: @@ -200,6 +201,11 @@ def execute_command( """ run_cmd = [] + # Note: the ssh user is + user = user or self.user() + if user: + run_cmd.extend(["sudo", "-E", "-u", user]) + use_sudo = kwargs.pop("sudo", self.use_sudo()) if use_sudo: run_cmd.extend(["sudo", "-E"])