xref: /dpdk/dts/framework/remote_session/interactive_remote_session.py (revision e9fd1ebf981f361844aea9ec94e17f4bda5e1479) 
1# SPDX-License-Identifier: BSD-3-Clause
2# Copyright(c) 2023 University of New Hampshire
3
4"""Handler for an SSH session dedicated to interactive shells."""
5
6import socket
7import traceback
8
9from paramiko import AutoAddPolicy, SSHClient, Transport  # type: ignore[import]
10from paramiko.ssh_exception import (  # type: ignore[import]
11    AuthenticationException,
12    BadHostKeyException,
13    NoValidConnectionsError,
14    SSHException,
15)
16
17from framework.config import NodeConfiguration
18from framework.exception import SSHConnectionError
19from framework.logger import DTSLOG
20
21
22class InteractiveRemoteSession:
23    """SSH connection dedicated to interactive applications.
24
25    The connection is created using `paramiko <https://docs.paramiko.org/en/latest/>`_
26    and is a persistent connection to the host. This class defines the methods for connecting
27    to the node and configures the connection to send "keep alive" packets every 30 seconds.
28    Because paramiko attempts to use SSH keys to establish a connection first, providing
29    a password is optional. This session is utilized by InteractiveShells
30    and cannot be interacted with directly.
31
32    Attributes:
33        hostname: The hostname that will be used to initialize a connection to the node.
34        ip: A subsection of `hostname` that removes the port for the connection if there
35            is one. If there is no port, this will be the same as hostname.
36        port: Port to use for the ssh connection. This will be extracted from `hostname`
37            if there is a port included, otherwise it will default to ``22``.
38        username: User to connect to the node with.
39        password: Password of the user connecting to the host. This will default to an
40            empty string if a password is not provided.
41        session: The underlying paramiko connection.
42
43    Raises:
44        SSHConnectionError: There is an error creating the SSH connection.
45    """
46
47    hostname: str
48    ip: str
49    port: int
50    username: str
51    password: str
52    session: SSHClient
53    _logger: DTSLOG
54    _node_config: NodeConfiguration
55    _transport: Transport | None
56
57    def __init__(self, node_config: NodeConfiguration, logger: DTSLOG) -> None:
58        """Connect to the node during initialization.
59
60        Args:
61            node_config: The test run configuration of the node to connect to.
62            logger: The logger instance this session will use.
63        """
64        self._node_config = node_config
65        self._logger = logger
66        self.hostname = node_config.hostname
67        self.username = node_config.user
68        self.password = node_config.password if node_config.password else ""
69        port = "22"
70        self.ip = node_config.hostname
71        if ":" in node_config.hostname:
72            self.ip, port = node_config.hostname.split(":")
73        self.port = int(port)
74        self._logger.info(
75            f"Initializing interactive connection for {self.username}@{self.hostname}"
76        )
77        self._connect()
78        self._logger.info(f"Interactive connection successful for {self.username}@{self.hostname}")
79
80    def _connect(self) -> None:
81        """Establish a connection to the node.
82
83        Connection attempts can be retried up to 10 times if certain errors are
84        encountered (NoValidConnectionsError, socket.error, SSHException). If a
85        connection is made, a 30 second "keep alive" interval will be set on the
86        session.
87
88        Raises:
89            SSHConnectionError: Connection cannot be established.
90        """
91        client = SSHClient()
92        client.set_missing_host_key_policy(AutoAddPolicy)
93        self.session = client
94        retry_attempts = 10
95        for retry_attempt in range(retry_attempts):
96            try:
97                client.connect(
98                    self.ip,
99                    username=self.username,
100                    port=self.port,
101                    password=self.password,
102                    timeout=20 if self.port else 10,
103                )
104            except (TypeError, BadHostKeyException, AuthenticationException) as e:
105                self._logger.exception(e)
106                raise SSHConnectionError(self.hostname) from e
107            except (NoValidConnectionsError, socket.error, SSHException) as e:
108                self._logger.debug(traceback.format_exc())
109                self._logger.warning(e)
110                self._logger.info(
111                    f"Retrying interactive session connection: retry number {retry_attempt +1}"
112                )
113            else:
114                break
115        else:
116            raise SSHConnectionError(self.hostname)
117        # Interactive sessions are used on an "as needed" basis so we have
118        # to set a keepalive
119        self._transport = self.session.get_transport()
120        if self._transport is not None:
121            self._transport.set_keepalive(30)
122