1# SPDX-License-Identifier: BSD-3-Clause 2# Copyright(c) 2023 University of New Hampshire 3 4"""Common functionality for interactive shell handling. 5 6The base class, :class:`InteractiveShell`, is meant to be extended by subclasses that contain 7functionality 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 13The :option:`--timeout` command line argument and the :envvar:`DTS_TIMEOUT` 14environment variable configure the timeout of getting the output from command execution. 15""" 16 17from abc import ABC 18from pathlib import PurePath 19from typing import Callable, ClassVar 20 21from paramiko import Channel, SSHClient, channel # type: ignore[import] 22 23from framework.logger import DTSLOG 24from framework.settings import SETTINGS 25 26 27class InteractiveShell(ABC): 28 """The base class for managing interactive shells. 29 30 This class shouldn't be instantiated directly, but instead be extended. It contains 31 methods for starting interactive shells as well as sending commands to these shells 32 and collecting input until reaching a certain prompt. All interactive applications 33 will use the same SSH connection, but each will create their own channel on that 34 session. 35 """ 36 37 _interactive_session: SSHClient 38 _stdin: channel.ChannelStdinFile 39 _stdout: channel.ChannelFile 40 _ssh_channel: Channel 41 _logger: DTSLOG 42 _timeout: float 43 _app_args: str 44 45 #: Prompt to expect at the end of output when sending a command. 46 #: This is often overridden by subclasses. 47 _default_prompt: ClassVar[str] = "" 48 49 #: Extra characters to add to the end of every command 50 #: before sending them. This is often overridden by subclasses and is 51 #: most commonly an additional newline character. 52 _command_extra_chars: ClassVar[str] = "" 53 54 #: Path to the executable to start the interactive application. 55 path: ClassVar[PurePath] 56 57 #: Whether this application is a DPDK app. If it is, the build directory 58 #: for DPDK on the node will be prepended to the path to the executable. 59 dpdk_app: ClassVar[bool] = False 60 61 def __init__( 62 self, 63 interactive_session: SSHClient, 64 logger: DTSLOG, 65 get_privileged_command: Callable[[str], str] | None, 66 app_args: str = "", 67 timeout: float = SETTINGS.timeout, 68 ) -> None: 69 """Create an SSH channel during initialization. 70 71 Args: 72 interactive_session: The SSH session dedicated to interactive shells. 73 logger: The logger instance this session will use. 74 get_privileged_command: A method for modifying a command to allow it to use 75 elevated privileges. If :data:`None`, the application will not be started 76 with elevated privileges. 77 app_args: The command line arguments to be passed to the application on startup. 78 timeout: The timeout used for the SSH channel that is dedicated to this interactive 79 shell. This timeout is for collecting output, so if reading from the buffer 80 and no output is gathered within the timeout, an exception is thrown. 81 """ 82 self._interactive_session = interactive_session 83 self._ssh_channel = self._interactive_session.invoke_shell() 84 self._stdin = self._ssh_channel.makefile_stdin("w") 85 self._stdout = self._ssh_channel.makefile("r") 86 self._ssh_channel.settimeout(timeout) 87 self._ssh_channel.set_combine_stderr(True) # combines stdout and stderr streams 88 self._logger = logger 89 self._timeout = timeout 90 self._app_args = app_args 91 self._start_application(get_privileged_command) 92 93 def _start_application(self, get_privileged_command: Callable[[str], str] | None) -> None: 94 """Starts a new interactive application based on the path to the app. 95 96 This method is often overridden by subclasses as their process for 97 starting may look different. 98 99 Args: 100 get_privileged_command: A function (but could be any callable) that produces 101 the version of the command with elevated privileges. 102 """ 103 start_command = f"{self.path} {self._app_args}" 104 if get_privileged_command is not None: 105 start_command = get_privileged_command(start_command) 106 self.send_command(start_command) 107 108 def send_command(self, command: str, prompt: str | None = None) -> str: 109 """Send `command` and get all output before the expected ending string. 110 111 Lines that expect input are not included in the stdout buffer, so they cannot 112 be used for expect. 113 114 Example: 115 If you were prompted to log into something with a username and password, 116 you cannot expect ``username:`` because it won't yet be in the stdout buffer. 117 A workaround for this could be consuming an extra newline character to force 118 the current `prompt` into the stdout buffer. 119 120 Args: 121 command: The command to send. 122 prompt: After sending the command, `send_command` will be expecting this string. 123 If :data:`None`, will use the class's default prompt. 124 125 Returns: 126 All output in the buffer before expected string. 127 """ 128 self._logger.info(f"Sending: '{command}'") 129 if prompt is None: 130 prompt = self._default_prompt 131 self._stdin.write(f"{command}{self._command_extra_chars}\n") 132 self._stdin.flush() 133 out: str = "" 134 for line in self._stdout: 135 out += line 136 if prompt in line and not line.rstrip().endswith( 137 command.rstrip() 138 ): # ignore line that sent command 139 break 140 self._logger.debug(f"Got output: {out}") 141 return out 142 143 def close(self) -> None: 144 """Properly free all resources.""" 145 self._stdin.close() 146 self._ssh_channel.close() 147 148 def __del__(self) -> None: 149 """Make sure the session is properly closed before deleting the object.""" 150 self.close() 151