frostfs-testcases/pytest_tests/helpers/ssh_helper.py

import logging
import socket
import tempfile
import textwrap
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from functools import wraps
from time import sleep
from typing import ClassVar, Optional

import allure
from paramiko import AutoAddPolicy, SFTPClient, SSHClient, SSHException, ssh_exception, RSAKey
from paramiko.ssh_exception import AuthenticationException


class HostIsNotAvailable(Exception):
    """Raises when host is not reachable."""

    def __init__(self, ip: str = None, exc: Exception = None):
        msg = f'Host is not available{f" by ip: {ip}" if ip else ""}'
        if exc:
            msg = f'{msg}. {exc}'
        super().__init__(msg)


def log_command(func):
    @wraps(func)
    def wrapper(host: 'HostClient', command: str, *args, **kwargs):
        display_length = 60
        short = command.removeprefix("$ProgressPreference='SilentlyContinue'\n")
        short = short[:display_length]
        short += '...' if short != command else ''
        with allure.step(f'SSH: {short}'):
            logging.info(f'Execute command "{command}" on "{host.ip}"')

            start_time = datetime.utcnow()
            cmd_result = func(host, command, *args, **kwargs)
            end_time = datetime.utcnow()

            log_message = f'HOST: {host.ip}\n' \
                          f'COMMAND:\n{textwrap.indent(command, " ")}\n' \
                          f'RC:\n {cmd_result.rc}\n' \
                          f'STDOUT:\n{textwrap.indent(cmd_result.stdout, " ")}\n' \
                          f'STDERR:\n{textwrap.indent(cmd_result.stderr, " ")}\n' \
                          f'Start / End / Elapsed\t {start_time.time()} / {end_time.time()} / {end_time - start_time}'

            logging.info(log_message)
            allure.attach(log_message, 'SSH command', allure.attachment_type.TEXT)
        return cmd_result
    return wrapper


@dataclass
class SSHCommand:
    stdout: str
    stderr: str
    rc: int


class HostClient:
    ssh_client: SSHClient
    SSH_CONNECTION_ATTEMPTS: ClassVar[int] = 3
    CONNECTION_TIMEOUT = 90

    TIMEOUT_RESTORE_CONNECTION = 10, 24

    def __init__(self, ip: str, login: str, password: Optional[str] = None,
                 private_key_path: Optional[str] = None, private_key_passphrase: Optional[str] = None,
                 init_ssh_client=True) -> None:
        self.ip = ip
        self.login = login
        self.password = password
        self.private_key_path = private_key_path
        self.private_key_passphrase = private_key_passphrase
        if init_ssh_client:
            self.create_connection(self.SSH_CONNECTION_ATTEMPTS)

    def exec(self, cmd: str, verify=True, timeout=90) -> SSHCommand:
        cmd_result = self._inner_exec(cmd, timeout)
        if verify:
            assert cmd_result.rc == 0, f'Non zero rc from command: "{cmd}"'
        return cmd_result

    @log_command
    def exec_with_confirmation(self, cmd: str, confirmation: list, verify=True, timeout=90) -> SSHCommand:
        ssh_stdin, ssh_stdout, ssh_stderr = self.ssh_client.exec_command(cmd, timeout=timeout)
        for line in confirmation:
            if not line.endswith('\n'):
                line = f'{line}\n'
            try:
                ssh_stdin.write(line)
            except OSError as err:
                logging.error(f'Got error {err} executing command {cmd}')
        ssh_stdin.close()
        output = SSHCommand(stdout=ssh_stdout.read().decode(errors='ignore'),
                            stderr=ssh_stderr.read().decode(errors='ignore'),
                            rc=ssh_stdout.channel.recv_exit_status())
        if verify:
            debug_info = f'\nSTDOUT: {output.stdout}\nSTDERR: {output.stderr}\nRC: {output.rc}'
            assert output.rc == 0, f'Non zero rc from command: "{cmd}"{debug_info}'
        return output

    @contextmanager
    def as_user(self, user: str, password: str):
        keep_user, keep_password = self.login, self.password
        self.login, self.password = user, password
        self.create_connection()
        yield
        self.login, self.password = keep_user, keep_password
        self.create_connection()

    @contextmanager
    def create_ssh_connection(self) -> 'SSHClient':
        if not self.ssh_client:
            self.create_connection()
        try:
            yield self.ssh_client
        finally:
            self.drop()

    @allure.step('Restore connection')
    def restore_ssh_connection(self):
        retry_time, retry_count = self.TIMEOUT_RESTORE_CONNECTION
        for _ in range(retry_count):
            try:
                self.create_connection()
            except AssertionError:
                logging.warning(f'Host: Cant reach host: {self.ip}.')
                sleep(retry_time)
            else:
                logging.info(f'Host: Cant reach host: {self.ip}.')
                return
        raise AssertionError(f'Host: Cant reach host: {self.ip} after 240 seconds..')

    @allure.step('Copy file {host_path_to_file} to local file {path_to_file}')
    def copy_file_from_host(self, host_path_to_file: str, path_to_file: str):
        with self._sftp_client() as sftp_client:
            sftp_client.get(host_path_to_file, path_to_file)

    def copy_file_to_host(self, path_to_file: str, host_path_to_file: str):
        with allure.step(f'Copy local file {path_to_file} to remote file {host_path_to_file} on host {self.ip}'):
            with self._sftp_client() as sftp_client:
                sftp_client.put(path_to_file, host_path_to_file)

    @allure.step('Save string to remote file {host_path_to_file}')
    def copy_str_to_host_file(self, string: str, host_path_to_file: str):
        with tempfile.NamedTemporaryFile(mode='r+') as temp:
            temp.writelines(string)
            temp.flush()
            with self._sftp_client() as client:
                client.put(temp.name, host_path_to_file)
        self.exec(f'cat {host_path_to_file}', verify=False)

    def create_connection(self, attempts=SSH_CONNECTION_ATTEMPTS):
        exc_err = None
        for attempt in range(attempts):
            self.ssh_client = SSHClient()
            self.ssh_client.set_missing_host_key_policy(AutoAddPolicy())
            try:
                if self.private_key_path:
                    logging.info(
                        f"Trying to connect to host {self.ip} as {self.login} using SSH key "
                        f"{self.private_key_path} (attempt {attempt})"
                    )
                    self.ssh_client.connect(
                        hostname=self.ip,
                        username=self.login,
                        pkey=RSAKey.from_private_key_file(self.private_key_path, self.private_key_passphrase),
                        timeout=self.CONNECTION_TIMEOUT
                    )
                else:
                    logging.info(
                        f"Trying to connect to host {self.ip} as {self.login} using password "
                        f"(attempt {attempt})"
                    )
                    self.ssh_client.connect(
                        hostname=self.ip,
                        username=self.login,
                        password=self.password,
                        timeout=self.CONNECTION_TIMEOUT
                    )
                return True

            except AuthenticationException as auth_err:
                logging.error(f'Host: {self.ip}. {auth_err}')
                self.drop()
                raise auth_err

            except (
                    SSHException,
                    ssh_exception.NoValidConnectionsError,
                    AttributeError,
                    socket.timeout,
                    OSError
            ) as ssh_err:
                exc_err = ssh_err
                self.drop()
                logging.error(f'Host: {self.ip}, connection error. {exc_err}')

        raise HostIsNotAvailable(self.ip, exc_err)

    def drop(self):
        self.ssh_client.close()

    @log_command
    def _inner_exec(self, cmd: str, timeout: int) -> SSHCommand:
        if not self.ssh_client:
            self.create_connection()
        for _ in range(self.SSH_CONNECTION_ATTEMPTS):
            try:
                _, stdout, stderr = self.ssh_client.exec_command(cmd, timeout=timeout)
                return SSHCommand(
                    stdout=stdout.read().decode(errors='ignore'),
                    stderr=stderr.read().decode(errors='ignore'),
                    rc=stdout.channel.recv_exit_status()
                )
            except (
                    SSHException,
                    TimeoutError,
                    ssh_exception.NoValidConnectionsError,
                    ConnectionResetError,
                    AttributeError,
                    socket.timeout,
            ) as ssh_err:
                logging.error(f'Host: {self.ip}, exec command error {ssh_err}')
                self.create_connection()
        raise HostIsNotAvailable(f'Host: {self.ip} is not reachable.')

    @contextmanager
    def _sftp_client(self) -> SFTPClient:
        with self.ssh_client.open_sftp() as sftp:
            yield sftp