1# SPDX-License-Identifier: BSD-3-Clause 2# Copyright(c) 2024 University of New Hampshire 3 4"""Common functionality for interactive shell handling. 5 6The base class, :class:`SingleActiveInteractiveShell`, is meant to be extended by subclasses that 7contain functionality specific to that shell type. These subclasses will often modify things like 8the prompt to expect or the arguments to pass into the application, but still utilize 9the same method for sending a command and collecting output. How this output is handled however 10is often application specific. If an application needs elevated privileges to start it is expected 11that the method for gaining those privileges is provided when initializing the class. 12 13This class is designed for applications like primary applications in DPDK where only one instance 14of the application can be running at a given time and, for this reason, is managed using a context 15manager. This context manager starts the application when you enter the context and cleans up the 16application when you exit. Using a context manager for this is useful since it allows us to ensure 17the application is cleaned up as soon as you leave the block regardless of the reason. 18 19The :option:`--timeout` command line argument and the :envvar:`DTS_TIMEOUT` 20environment variable configure the timeout of getting the output from command execution. 21""" 22 23from abc import ABC 24from pathlib import PurePath 25from typing import ClassVar 26 27from paramiko import Channel, channel # type: ignore[import-untyped] 28from typing_extensions import Self 29 30from framework.exception import ( 31 InteractiveCommandExecutionError, 32 InteractiveSSHSessionDeadError, 33 InteractiveSSHTimeoutError, 34) 35from framework.logger import DTSLogger, get_dts_logger 36from framework.params import Params 37from framework.settings import SETTINGS 38from framework.testbed_model.node import Node 39from framework.utils import MultiInheritanceBaseClass 40 41 42class SingleActiveInteractiveShell(MultiInheritanceBaseClass, ABC): 43 """The base class for managing interactive shells. 44 45 This class shouldn't be instantiated directly, but instead be extended. It contains 46 methods for starting interactive shells as well as sending commands to these shells 47 and collecting input until reaching a certain prompt. All interactive applications 48 will use the same SSH connection, but each will create their own channel on that 49 session. 50 51 Interactive shells are started and stopped using a context manager. This allows for the start 52 and cleanup of the application to happen at predictable times regardless of exceptions or 53 interrupts. 54 55 Attributes: 56 is_alive: :data:`True` if the application has started successfully, :data:`False` 57 otherwise. 58 """ 59 60 _node: Node 61 _stdin: channel.ChannelStdinFile 62 _stdout: channel.ChannelFile 63 _ssh_channel: Channel 64 _logger: DTSLogger 65 _timeout: float 66 _app_params: Params 67 _privileged: bool 68 _real_path: PurePath 69 70 #: The number of times to try starting the application before considering it a failure. 71 _init_attempts: ClassVar[int] = 5 72 73 #: Prompt to expect at the end of output when sending a command. 74 #: This is often overridden by subclasses. 75 _default_prompt: ClassVar[str] = "" 76 77 #: Extra characters to add to the end of every command 78 #: before sending them. This is often overridden by subclasses and is 79 #: most commonly an additional newline character. This additional newline 80 #: character is used to force the line that is currently awaiting input 81 #: into the stdout buffer so that it can be consumed and checked against 82 #: the expected prompt. 83 _command_extra_chars: ClassVar[str] = "" 84 85 #: Path to the executable to start the interactive application. 86 path: ClassVar[PurePath] 87 88 is_alive: bool = False 89 90 def __init__( 91 self, 92 node: Node, 93 privileged: bool = False, 94 timeout: float = SETTINGS.timeout, 95 app_params: Params = Params(), 96 name: str | None = None, 97 **kwargs, 98 ) -> None: 99 """Create an SSH channel during initialization. 100 101 Additional keyword arguments can be passed through `kwargs` if needed for fulfilling other 102 constructors in the case of multiple inheritance. 103 104 Args: 105 node: The node on which to run start the interactive shell. 106 privileged: Enables the shell to run as superuser. 107 timeout: The timeout used for the SSH channel that is dedicated to this interactive 108 shell. This timeout is for collecting output, so if reading from the buffer 109 and no output is gathered within the timeout, an exception is thrown. 110 app_params: The command line parameters to be passed to the application on startup. 111 name: Name for the interactive shell to use for logging. This name will be appended to 112 the name of the underlying node which it is running on. 113 """ 114 self._node = node 115 if name is None: 116 name = type(self).__name__ 117 self._logger = get_dts_logger(f"{node.name}.{name}") 118 self._app_params = app_params 119 self._privileged = privileged 120 self._timeout = timeout 121 # Ensure path is properly formatted for the host 122 self._update_real_path(self.path) 123 super().__init__(node, **kwargs) 124 125 def _setup_ssh_channel(self): 126 self._ssh_channel = self._node.main_session.interactive_session.session.invoke_shell() 127 self._stdin = self._ssh_channel.makefile_stdin("w") 128 self._stdout = self._ssh_channel.makefile("r") 129 self._ssh_channel.settimeout(self._timeout) 130 self._ssh_channel.set_combine_stderr(True) # combines stdout and stderr streams 131 132 def _make_start_command(self) -> str: 133 """Makes the command that starts the interactive shell.""" 134 start_command = f"{self._real_path} {self._app_params or ''}" 135 if self._privileged: 136 start_command = self._node.main_session._get_privileged_command(start_command) 137 return start_command 138 139 def _start_application(self) -> None: 140 """Starts a new interactive application based on the path to the app. 141 142 This method is often overridden by subclasses as their process for starting may look 143 different. Initialization of the shell on the host can be retried up to 144 `self._init_attempts` - 1 times. This is done because some DPDK applications need slightly 145 more time after exiting their script to clean up EAL before others can start. 146 147 Raises: 148 InteractiveCommandExecutionError: If the application fails to start within the allotted 149 number of retries. 150 """ 151 self._setup_ssh_channel() 152 self._ssh_channel.settimeout(5) 153 start_command = self._make_start_command() 154 self.is_alive = True 155 for attempt in range(self._init_attempts): 156 try: 157 self.send_command(start_command) 158 break 159 except InteractiveSSHTimeoutError: 160 self._logger.info( 161 f"Interactive shell failed to start (attempt {attempt+1} out of " 162 f"{self._init_attempts})" 163 ) 164 else: 165 self._ssh_channel.settimeout(self._timeout) 166 self.is_alive = False # update state on failure to start 167 raise InteractiveCommandExecutionError("Failed to start application.") 168 self._ssh_channel.settimeout(self._timeout) 169 170 def send_command( 171 self, command: str, prompt: str | None = None, skip_first_line: bool = False 172 ) -> str: 173 """Send `command` and get all output before the expected ending string. 174 175 Lines that expect input are not included in the stdout buffer, so they cannot 176 be used for expect. 177 178 Example: 179 If you were prompted to log into something with a username and password, 180 you cannot expect ``username:`` because it won't yet be in the stdout buffer. 181 A workaround for this could be consuming an extra newline character to force 182 the current `prompt` into the stdout buffer. 183 184 Args: 185 command: The command to send. 186 prompt: After sending the command, `send_command` will be expecting this string. 187 If :data:`None`, will use the class's default prompt. 188 skip_first_line: Skip the first line when capturing the output. 189 190 Returns: 191 All output in the buffer before expected string. 192 193 Raises: 194 InteractiveCommandExecutionError: If attempting to send a command to a shell that is 195 not currently running. 196 InteractiveSSHSessionDeadError: The session died while executing the command. 197 InteractiveSSHTimeoutError: If command was sent but prompt could not be found in 198 the output before the timeout. 199 """ 200 if not self.is_alive: 201 raise InteractiveCommandExecutionError( 202 f"Cannot send command {command} to application because the shell is not running." 203 ) 204 self._logger.info(f"Sending: '{command}'") 205 if prompt is None: 206 prompt = self._default_prompt 207 out: str = "" 208 try: 209 self._stdin.write(f"{command}{self._command_extra_chars}\n") 210 self._stdin.flush() 211 for line in self._stdout: 212 if skip_first_line: 213 skip_first_line = False 214 continue 215 if line.rstrip().endswith(prompt): 216 break 217 out += line 218 except TimeoutError as e: 219 self._logger.exception(e) 220 self._logger.debug( 221 f"Prompt ({prompt}) was not found in output from command before timeout." 222 ) 223 raise InteractiveSSHTimeoutError(command) from e 224 except OSError as e: 225 self._logger.exception(e) 226 raise InteractiveSSHSessionDeadError( 227 self._node.main_session.interactive_session.hostname 228 ) from e 229 finally: 230 self._logger.debug(f"Got output: {out}") 231 return out 232 233 def _close(self) -> None: 234 self._stdin.close() 235 self._ssh_channel.close() 236 237 def _update_real_path(self, path: PurePath) -> None: 238 """Updates the interactive shell's real path used at command line.""" 239 self._real_path = self._node.main_session.join_remote_path(path) 240 241 def __enter__(self) -> Self: 242 """Enter the context block. 243 244 Upon entering a context block with this class, the desired behavior is to create the 245 channel for the application to use, and then start the application. 246 247 Returns: 248 Reference to the object for the application after it has been started. 249 """ 250 self._start_application() 251 return self 252 253 def __exit__(self, *_) -> None: 254 """Exit the context block. 255 256 Upon exiting a context block with this class, we want to ensure that the instance of the 257 application is explicitly closed and properly cleaned up using its close method. Note that 258 because this method returns :data:`None` if an exception was raised within the block, it is 259 not handled and will be re-raised after the application is closed. 260 261 The desired behavior is to close the application regardless of the reason for exiting the 262 context and then recreate that reason afterwards. All method arguments are ignored for 263 this reason. 264 """ 265 self._close() 266