[v2,1/3] dts: Improve output gathering in interactive shells

diff --git a/dts/framework/exception.py b/dts/framework/exception.py
index cce1e0231a..627190c781 100644
--- a/dts/framework/exception.py
+++ b/dts/framework/exception.py
@@ -48,26 +48,6 @@  class DTSError(Exception):
     severity: ClassVar[ErrorSeverity] = ErrorSeverity.GENERIC_ERR
 
 
-class SSHTimeoutError(DTSError):
-    """The SSH execution of a command timed out."""
-
-    #:
-    severity: ClassVar[ErrorSeverity] = ErrorSeverity.SSH_ERR
-    _command: str
-
-    def __init__(self, command: str):
-        """Define the meaning of the first argument.
-
-        Args:
-            command: The executed command.
-        """
-        self._command = command
-
-    def __str__(self) -> str:
-        """Add some context to the string representation."""
-        return f"{self._command} execution timed out."
-
-
 class SSHConnectionError(DTSError):
     """An unsuccessful SSH connection."""
 
@@ -95,8 +75,42 @@  def __str__(self) -> str:
         return message
 
 
-class SSHSessionDeadError(DTSError):
-    """The SSH session is no longer alive."""
+class _SSHTimeoutError(DTSError):
+    """The execution of a command via SSH timed out.
+
+    This class is private and meant to be raised as its interactive and non-interactive variants.
+    """
+
+    #:
+    severity: ClassVar[ErrorSeverity] = ErrorSeverity.SSH_ERR
+    _command: str
+
+    def __init__(self, command: str):
+        """Define the meaning of the first argument.
+
+        Args:
+            command: The executed command.
+        """
+        self._command = command
+
+    def __str__(self) -> str:
+        """Add some context to the string representation."""
+        return f"{self._command} execution timed out."
+
+
+class SSHTimeoutError(_SSHTimeoutError):
+    """The execution of a command on a non-interactive SSH session timed out."""
+
+
+class InteractiveSSHTimeoutError(_SSHTimeoutError):
+    """The execution of a command on an interactive SSH session timed out."""
+
+
+class _SSHSessionDeadError(DTSError):
+    """The SSH session is no longer alive.
+
+    This class is private and meant to be raised as its interactive and non-interactive variants.
+    """
 
     #:
     severity: ClassVar[ErrorSeverity] = ErrorSeverity.SSH_ERR
@@ -115,6 +129,14 @@  def __str__(self) -> str:
         return f"SSH session with {self._host} has died."
 
 
+class SSHSessionDeadError(_SSHSessionDeadError):
+    """Non-interactive SSH session has died."""
+
+
+class InteractiveSSHSessionDeadError(_SSHSessionDeadError):
+    """Interactive SSH session as died."""
+
+
 class ConfigurationError(DTSError):
     """An invalid configuration."""
 
diff --git a/dts/framework/remote_session/interactive_shell.py b/dts/framework/remote_session/interactive_shell.py
index 5cfe202e15..0b0ccdb545 100644
--- a/dts/framework/remote_session/interactive_shell.py
+++ b/dts/framework/remote_session/interactive_shell.py
@@ -18,11 +18,17 @@ 
 from pathlib import PurePath
 from typing import Callable, ClassVar
 
-from paramiko import Channel, SSHClient, channel  # type: ignore[import]
+from paramiko import Channel, channel  # type: ignore[import]
 
+from framework.exception import (
+    InteractiveSSHSessionDeadError,
+    InteractiveSSHTimeoutError,
+)
 from framework.logger import DTSLogger
 from framework.settings import SETTINGS
 
+from .interactive_remote_session import InteractiveRemoteSession
+
 
 class InteractiveShell(ABC):
     """The base class for managing interactive shells.
@@ -34,7 +40,7 @@  class InteractiveShell(ABC):
     session.
     """
 
-    _interactive_session: SSHClient
+    _interactive_session: InteractiveRemoteSession
     _stdin: channel.ChannelStdinFile
     _stdout: channel.ChannelFile
     _ssh_channel: Channel
@@ -60,7 +66,7 @@  class InteractiveShell(ABC):
 
     def __init__(
         self,
-        interactive_session: SSHClient,
+        interactive_session: InteractiveRemoteSession,
         logger: DTSLogger,
         get_privileged_command: Callable[[str], str] | None,
         app_args: str = "",
@@ -80,7 +86,7 @@  def __init__(
                 and no output is gathered within the timeout, an exception is thrown.
         """
         self._interactive_session = interactive_session
-        self._ssh_channel = self._interactive_session.invoke_shell()
+        self._ssh_channel = self._interactive_session.session.invoke_shell()
         self._stdin = self._ssh_channel.makefile_stdin("w")
         self._stdout = self._ssh_channel.makefile("r")
         self._ssh_channel.settimeout(timeout)
@@ -124,20 +130,34 @@  def send_command(self, command: str, prompt: str | None = None) -> str:
 
         Returns:
             All output in the buffer before expected string.
+
+        Raises:
+            InteractiveSSHSessionDeadError: The session died while executing the command.
+            InteractiveSSHTimeoutError: If command was sent but prompt could not be found in
+                the output before the timeout.
         """
         self._logger.info(f"Sending: '{command}'")
         if prompt is None:
             prompt = self._default_prompt
-        self._stdin.write(f"{command}{self._command_extra_chars}\n")
-        self._stdin.flush()
         out: str = ""
-        for line in self._stdout:
-            out += line
-            if prompt in line and not line.rstrip().endswith(
-                command.rstrip()
-            ):  # ignore line that sent command
-                break
-        self._logger.debug(f"Got output: {out}")
+        try:
+            self._stdin.write(f"{command}{self._command_extra_chars}\n")
+            self._stdin.flush()
+            for line in self._stdout:
+                out += line
+                if line.rstrip().endswith(prompt):
+                    break
+        except TimeoutError as e:
+            self._logger.exception(e)
+            self._logger.debug(
+                f"Prompt ({prompt}) was not found in output from command before timeout."
+            )
+            raise InteractiveSSHTimeoutError(command) from e
+        except OSError as e:
+            self._logger.exception(e)
+            raise InteractiveSSHSessionDeadError(self._interactive_session.hostname) from e
+        finally:
+            self._logger.debug(f"Got output: {out}")
         return out
 
     def close(self) -> None: