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