143 lines
5.4 KiB
Python
143 lines
5.4 KiB
Python
#
|
|
# Copyright(c) 2019-2021 Intel Corporation
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
#
|
|
|
|
import socket
|
|
import subprocess
|
|
import paramiko
|
|
|
|
from datetime import timedelta, datetime
|
|
from connection.base_executor import BaseExecutor
|
|
from core.test_run import TestRun
|
|
from test_utils.output import Output
|
|
|
|
|
|
class SshExecutor(BaseExecutor):
|
|
def __init__(self, ip, username, port=22):
|
|
self.ip = ip
|
|
self.user = username
|
|
self.port = port
|
|
self.ssh = paramiko.SSHClient()
|
|
self._check_config_for_reboot_timeout()
|
|
|
|
def __del__(self):
|
|
self.ssh.close()
|
|
|
|
def connect(self, user=None, port=None,
|
|
timeout: timedelta = timedelta(seconds=30)):
|
|
user = user or self.user
|
|
port = port or self.port
|
|
self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
try:
|
|
self.ssh.connect(self.ip, username=user,
|
|
port=port, timeout=timeout.total_seconds(),
|
|
banner_timeout=timeout.total_seconds())
|
|
except paramiko.AuthenticationException as e:
|
|
raise paramiko.AuthenticationException(
|
|
f"Authentication exception occurred while trying to connect to DUT. "
|
|
f"Please check your SSH key-based authentication.\n{e}")
|
|
except (paramiko.SSHException, socket.timeout) as e:
|
|
raise ConnectionError(f"An exception of type '{type(e)}' occurred while trying to "
|
|
f"connect to {self.ip}.\n {e}")
|
|
|
|
def disconnect(self):
|
|
try:
|
|
self.ssh.close()
|
|
except Exception:
|
|
raise Exception(f"An exception occurred while trying to disconnect from {self.ip}")
|
|
|
|
def _execute(self, command, timeout):
|
|
try:
|
|
(stdin, stdout, stderr) = self.ssh.exec_command(command,
|
|
timeout=timeout.total_seconds())
|
|
except paramiko.SSHException as e:
|
|
raise ConnectionError(f"An exception occurred while executing command '{command}' on"
|
|
f" {self.ip}\n{e}")
|
|
|
|
return Output(stdout.read(), stderr.read(), stdout.channel.recv_exit_status())
|
|
|
|
def _rsync(self, src, dst, delete=False, symlinks=False, checksum=False, exclude_list=[],
|
|
timeout: timedelta = timedelta(seconds=90), dut_to_controller=False):
|
|
options = []
|
|
|
|
if delete:
|
|
options.append("--delete")
|
|
if symlinks:
|
|
options.append("--links")
|
|
if checksum:
|
|
options.append("--checksum")
|
|
|
|
for exclude in exclude_list:
|
|
options.append(f"--exclude {exclude}")
|
|
|
|
src_to_dst = f"{self.user}@{self.ip}:{src} {dst} " if dut_to_controller else\
|
|
f"{src} {self.user}@{self.ip}:{dst} "
|
|
|
|
try:
|
|
completed_process = subprocess.run(
|
|
f'rsync -r -e "ssh -p {self.port} '
|
|
f'-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no" '
|
|
+ src_to_dst + f'{" ".join(options)}',
|
|
shell=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
timeout=timeout.total_seconds())
|
|
except Exception as e:
|
|
TestRun.LOGGER.exception(f"Exception occurred during rsync process. "
|
|
f"Please check your SSH key-based authentication.\n{e}")
|
|
|
|
if completed_process.returncode:
|
|
raise Exception(f"rsync failed:\n{completed_process}")
|
|
|
|
def is_remote(self):
|
|
return True
|
|
|
|
def _check_config_for_reboot_timeout(self):
|
|
if "reboot_timeout" in TestRun.config.keys():
|
|
self._parse_timeout_to_int()
|
|
else:
|
|
self.reboot_timeout = None
|
|
|
|
def _parse_timeout_to_int(self):
|
|
self.reboot_timeout = int(TestRun.config["reboot_timeout"])
|
|
if self.reboot_timeout < 0:
|
|
raise ValueError("Reboot timeout cannot be negative.")
|
|
|
|
def reboot(self):
|
|
self.run("reboot")
|
|
self.wait_for_connection_loss()
|
|
self.wait_for_connection(timedelta(seconds=self.reboot_timeout)) \
|
|
if self.reboot_timeout is not None else self.wait_for_connection()
|
|
|
|
def is_active(self):
|
|
try:
|
|
self.ssh.exec_command('', timeout=5)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def wait_for_connection(self, timeout: timedelta = timedelta(minutes=10)):
|
|
start_time = datetime.now()
|
|
with TestRun.group("Waiting for DUT ssh connection"):
|
|
while start_time + timeout > datetime.now():
|
|
try:
|
|
self.connect()
|
|
return
|
|
except paramiko.AuthenticationException:
|
|
raise
|
|
except Exception:
|
|
continue
|
|
raise ConnectionError("Timeout occurred while trying to establish ssh connection")
|
|
|
|
def wait_for_connection_loss(self, timeout: timedelta = timedelta(minutes=1)):
|
|
with TestRun.group("Waiting for DUT ssh connection loss"):
|
|
end_time = datetime.now() + timeout
|
|
while end_time > datetime.now():
|
|
self.disconnect()
|
|
try:
|
|
self.connect(timeout=timedelta(seconds=5))
|
|
except Exception:
|
|
return
|
|
raise ConnectionError("Timeout occurred before ssh connection loss")
|