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