path: root/src/pybind/mgr/cephadm/ssh.py

import logging
import os
import asyncio
import concurrent
from tempfile import NamedTemporaryFile
from threading import Thread
from contextlib import contextmanager
from io import StringIO
from shlex import quote
from typing import TYPE_CHECKING, Optional, List, Tuple, Dict, Iterator, TypeVar, Awaitable, Union
from orchestrator import OrchestratorError

try:
    import asyncssh
except ImportError:
    asyncssh = None  # type: ignore

if TYPE_CHECKING:
    from cephadm.module import CephadmOrchestrator
    from asyncssh.connection import SSHClientConnection

T = TypeVar('T')


logger = logging.getLogger(__name__)

asyncssh_logger = logging.getLogger('asyncssh')
asyncssh_logger.propagate = False


class HostConnectionError(OrchestratorError):
    def __init__(self, message: str, hostname: str, addr: str) -> None:
        super().__init__(message)
        self.hostname = hostname
        self.addr = addr


DEFAULT_SSH_CONFIG = """
Host *
  User root
  StrictHostKeyChecking no
  UserKnownHostsFile /dev/null
  ConnectTimeout=30
"""


class EventLoopThread(Thread):

    def __init__(self) -> None:
        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)

        super().__init__(target=self._loop.run_forever)
        self.start()

    def get_result(self, coro: Awaitable[T], timeout: Optional[int] = None) -> T:
        # useful to note: This "run_coroutine_threadsafe" returns a
        # concurrent.futures.Future, rather than an asyncio.Future. They are
        # fairly similar but have a few differences, notably in our case
        # that the result function of a concurrent.futures.Future accepts
        # a timeout argument
        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
        try:
            return future.result(timeout)
        except (asyncio.TimeoutError, concurrent.futures.TimeoutError):
            # try to cancel the task before raising the exception further up
            future.cancel()
            raise


class SSHManager:

    def __init__(self, mgr: "CephadmOrchestrator"):
        self.mgr: "CephadmOrchestrator" = mgr
        self.cons: Dict[str, "SSHClientConnection"] = {}

    async def _remote_connection(self,
                                 host: str,
                                 addr: Optional[str] = None,
                                 ) -> "SSHClientConnection":
        if not self.cons.get(host) or host not in self.mgr.inventory:
            if not addr and host in self.mgr.inventory:
                addr = self.mgr.inventory.get_addr(host)

            if not addr:
                raise OrchestratorError("host address is empty")

            assert self.mgr.ssh_user
            n = self.mgr.ssh_user + '@' + addr
            logger.debug("Opening connection to {} with ssh options '{}'".format(
                n, self.mgr._ssh_options))

            asyncssh.set_log_level('DEBUG')
            asyncssh.set_debug_level(3)

            with self.redirect_log(host, addr):
                try:
                    ssh_options = asyncssh.SSHClientConnectionOptions(
                        keepalive_interval=7, keepalive_count_max=3)
                    conn = await asyncssh.connect(addr, username=self.mgr.ssh_user, client_keys=[self.mgr.tkey.name],
                                                  known_hosts=None, config=[self.mgr.ssh_config_fname],
                                                  preferred_auth=['publickey'], options=ssh_options)
                except OSError:
                    raise
                except asyncssh.Error:
                    raise
                except Exception:
                    raise
            self.cons[host] = conn

        self.mgr.offline_hosts_remove(host)

        return self.cons[host]

    @contextmanager
    def redirect_log(self, host: str, addr: str) -> Iterator[None]:
        log_string = StringIO()
        ch = logging.StreamHandler(log_string)
        ch.setLevel(logging.INFO)
        asyncssh_logger.addHandler(ch)

        try:
            yield
        except OSError as e:
            self.mgr.offline_hosts.add(host)
            log_content = log_string.getvalue()
            msg = f"Can't communicate with remote host `{addr}`, possibly because the host is not reachable or python3 is not installed on the host. {str(e)}"
            logger.exception(msg)
            raise HostConnectionError(msg, host, addr)
        except asyncssh.Error as e:
            self.mgr.offline_hosts.add(host)
            log_content = log_string.getvalue()
            msg = f'Failed to connect to {host} ({addr}). {str(e)}' + '\n' + f'Log: {log_content}'
            logger.debug(msg)
            raise HostConnectionError(msg, host, addr)
        except Exception as e:
            self.mgr.offline_hosts.add(host)
            log_content = log_string.getvalue()
            logger.exception(str(e))
            raise HostConnectionError(
                f'Failed to connect to {host} ({addr}): {repr(e)}' + '\n' f'Log: {log_content}', host, addr)
        finally:
            log_string.flush()
            asyncssh_logger.removeHandler(ch)

    def remote_connection(self,
                          host: str,
                          addr: Optional[str] = None,
                          ) -> "SSHClientConnection":
        with self.mgr.async_timeout_handler(host, f'ssh {host} (addr {addr})'):
            return self.mgr.wait_async(self._remote_connection(host, addr))

    async def _execute_command(self,
                               host: str,
                               cmd_components: List[str],
                               stdin: Optional[str] = None,
                               addr: Optional[str] = None,
                               log_command: Optional[bool] = True,
                               ) -> Tuple[str, str, int]:

        conn = await self._remote_connection(host, addr)
        sudo_prefix = "sudo " if self.mgr.ssh_user != 'root' else ""
        cmd = sudo_prefix + " ".join(quote(x) for x in cmd_components)
        try:
            address = addr or self.mgr.inventory.get_addr(host)
        except Exception:
            address = host
        if log_command:
            logger.debug(f'Running command: {cmd}')
        try:
            r = await conn.run(f'{sudo_prefix}true', check=True, timeout=5)  # host quick check
            r = await conn.run(cmd, input=stdin)
        # handle these Exceptions otherwise you might get a weird error like
        # TypeError: __init__() missing 1 required positional argument: 'reason' (due to the asyncssh error interacting with raise_if_exception)
        except asyncssh.ChannelOpenError as e:
            # SSH connection closed or broken, will create new connection next call
            logger.debug(f'Connection to {host} failed. {str(e)}')
            await self._reset_con(host)
            self.mgr.offline_hosts.add(host)
            raise HostConnectionError(f'Unable to reach remote host {host}. {str(e)}', host, address)
        except asyncssh.ProcessError as e:
            msg = f"Cannot execute the command '{cmd}' on the {host}. {str(e.stderr)}."
            logger.debug(msg)
            await self._reset_con(host)
            self.mgr.offline_hosts.add(host)
            raise HostConnectionError(msg, host, address)
        except Exception as e:
            msg = f"Generic error while executing command '{cmd}' on the host {host}. {str(e)}."
            logger.debug(msg)
            await self._reset_con(host)
            self.mgr.offline_hosts.add(host)
            raise HostConnectionError(msg, host, address)

        def _rstrip(v: Union[bytes, str, None]) -> str:
            if not v:
                return ''
            if isinstance(v, str):
                return v.rstrip('\n')
            if isinstance(v, bytes):
                return v.decode().rstrip('\n')
            raise OrchestratorError(
                f'Unable to parse ssh output with type {type(v)} from remote host {host}')

        out = _rstrip(r.stdout)
        err = _rstrip(r.stderr)
        rc = r.returncode if r.returncode else 0

        return out, err, rc

    def execute_command(self,
                        host: str,
                        cmd: List[str],
                        stdin: Optional[str] = None,
                        addr: Optional[str] = None,
                        log_command: Optional[bool] = True
                        ) -> Tuple[str, str, int]:
        with self.mgr.async_timeout_handler(host, " ".join(cmd)):
            return self.mgr.wait_async(self._execute_command(host, cmd, stdin, addr, log_command))

    async def _check_execute_command(self,
                                     host: str,
                                     cmd: List[str],
                                     stdin: Optional[str] = None,
                                     addr: Optional[str] = None,
                                     log_command: Optional[bool] = True
                                     ) -> str:
        out, err, code = await self._execute_command(host, cmd, stdin, addr, log_command)
        if code != 0:
            msg = f'Command {cmd} failed. {err}'
            logger.debug(msg)
            raise OrchestratorError(msg)
        return out

    def check_execute_command(self,
                              host: str,
                              cmd: List[str],
                              stdin: Optional[str] = None,
                              addr: Optional[str] = None,
                              log_command: Optional[bool] = True,
                              ) -> str:
        with self.mgr.async_timeout_handler(host, " ".join(cmd)):
            return self.mgr.wait_async(self._check_execute_command(host, cmd, stdin, addr, log_command))

    async def _write_remote_file(self,
                                 host: str,
                                 path: str,
                                 content: bytes,
                                 mode: Optional[int] = None,
                                 uid: Optional[int] = None,
                                 gid: Optional[int] = None,
                                 addr: Optional[str] = None,
                                 ) -> None:
        try:
            cephadm_tmp_dir = f"/tmp/cephadm-{self.mgr._cluster_fsid}"
            dirname = os.path.dirname(path)
            await self._check_execute_command(host, ['mkdir', '-p', dirname], addr=addr)
            await self._check_execute_command(host, ['mkdir', '-p', cephadm_tmp_dir + dirname], addr=addr)
            tmp_path = cephadm_tmp_dir + path + '.new'
            await self._check_execute_command(host, ['touch', tmp_path], addr=addr)
            if self.mgr.ssh_user != 'root':
                assert self.mgr.ssh_user
                await self._check_execute_command(host, ['chown', '-R', self.mgr.ssh_user, cephadm_tmp_dir], addr=addr)
                await self._check_execute_command(host, ['chmod', str(644), tmp_path], addr=addr)
            with NamedTemporaryFile(prefix='cephadm-write-remote-file-') as f:
                os.fchmod(f.fileno(), 0o600)
                f.write(content)
                f.flush()
                conn = await self._remote_connection(host, addr)
                async with conn.start_sftp_client() as sftp:
                    await sftp.put(f.name, tmp_path)
            if uid is not None and gid is not None and mode is not None:
                # shlex quote takes str or byte object, not int
                await self._check_execute_command(host, ['chown', '-R', str(uid) + ':' + str(gid), tmp_path], addr=addr)
                await self._check_execute_command(host, ['chmod', oct(mode)[2:], tmp_path], addr=addr)
            await self._check_execute_command(host, ['mv', tmp_path, path], addr=addr)
        except Exception as e:
            msg = f"Unable to write {host}:{path}: {e}"
            logger.exception(msg)
            raise OrchestratorError(msg)

    def write_remote_file(self,
                          host: str,
                          path: str,
                          content: bytes,
                          mode: Optional[int] = None,
                          uid: Optional[int] = None,
                          gid: Optional[int] = None,
                          addr: Optional[str] = None,
                          ) -> None:
        with self.mgr.async_timeout_handler(host, f'writing file {path}'):
            self.mgr.wait_async(self._write_remote_file(
                host, path, content, mode, uid, gid, addr))

    async def _reset_con(self, host: str) -> None:
        conn = self.cons.get(host)
        if conn:
            logger.debug(f'_reset_con close {host}')
            conn.close()
            del self.cons[host]

    def reset_con(self, host: str) -> None:
        with self.mgr.async_timeout_handler(cmd=f'resetting ssh connection to {host}'):
            self.mgr.wait_async(self._reset_con(host))

    def _reset_cons(self) -> None:
        for host, conn in self.cons.items():
            logger.debug(f'_reset_cons close {host}')
            conn.close()
        self.cons = {}

    def _reconfig_ssh(self) -> None:
        temp_files = []  # type: list
        ssh_options = []  # type: List[str]

        # ssh_config
        self.mgr.ssh_config_fname = self.mgr.ssh_config_file
        ssh_config = self.mgr.get_store("ssh_config")
        if ssh_config is not None or self.mgr.ssh_config_fname is None:
            if not ssh_config:
                ssh_config = DEFAULT_SSH_CONFIG
            f = NamedTemporaryFile(prefix='cephadm-conf-')
            os.fchmod(f.fileno(), 0o600)
            f.write(ssh_config.encode('utf-8'))
            f.flush()  # make visible to other processes
            temp_files += [f]
            self.mgr.ssh_config_fname = f.name
        if self.mgr.ssh_config_fname:
            self.mgr.validate_ssh_config_fname(self.mgr.ssh_config_fname)
            ssh_options += ['-F', self.mgr.ssh_config_fname]
        self.mgr.ssh_config = ssh_config

        # identity
        ssh_key = self.mgr.get_store("ssh_identity_key")
        ssh_pub = self.mgr.get_store("ssh_identity_pub")
        ssh_cert = self.mgr.get_store("ssh_identity_cert")
        self.mgr.ssh_pub = ssh_pub
        self.mgr.ssh_key = ssh_key
        self.mgr.ssh_cert = ssh_cert
        if ssh_key:
            self.mgr.tkey = NamedTemporaryFile(prefix='cephadm-identity-')
            self.mgr.tkey.write(ssh_key.encode('utf-8'))
            os.fchmod(self.mgr.tkey.fileno(), 0o600)
            self.mgr.tkey.flush()  # make visible to other processes
            temp_files += [self.mgr.tkey]
            if ssh_pub:
                tpub = open(self.mgr.tkey.name + '.pub', 'w')
                os.fchmod(tpub.fileno(), 0o600)
                tpub.write(ssh_pub)
                tpub.flush()  # make visible to other processes
                temp_files += [tpub]
            if ssh_cert:
                tcert = open(self.mgr.tkey.name + '-cert.pub', 'w')
                os.fchmod(tcert.fileno(), 0o600)
                tcert.write(ssh_cert)
                tcert.flush()  # make visible to other processes
                temp_files += [tcert]
            ssh_options += ['-i', self.mgr.tkey.name]

        self.mgr._temp_files = temp_files
        if ssh_options:
            self.mgr._ssh_options = ' '.join(ssh_options)
        else:
            self.mgr._ssh_options = None

        if self.mgr.mode == 'root':
            self.mgr.ssh_user = self.mgr.get_store('ssh_user', default='root')
        elif self.mgr.mode == 'cephadm-package':
            self.mgr.ssh_user = 'cephadm'

        self._reset_cons()