from __future__ import annotations import abc import os import pwd import sys from enum import Flag, auto from typing import TYPE_CHECKING from ..base import Result from ..ExecContext import ExecContext from ..log import DEBUG, ERR, INFO, NOTICE, WARNING, log from ..Uri import Uri if TYPE_CHECKING: from ..Uri import Uri class SSHClient(ExecContext): class Caps(Flag): LogOutput = auto() Interactive = auto() ModEnv = auto() Wd = auto() def __init__(self, uri: Uri | str, caps: Caps = Caps(0), *args, **kwargs) -> None: uri = Uri.pimp(uri) if uri.username is None: uri.set_username(pwd.getpwuid(os.getuid()).pw_name) super().__init__(uri = uri, *args, **kwargs) self.__caps = caps @abc.abstractmethod async def _run_ssh( self, cmd: list[str], wd: str | None, verbose: bool, cmd_input: bytes | None, mod_env: dict[str, str] | None, interactive: bool, log_prefix: str, ) -> Result: pass async def _run( self, cmd: list[str], wd: str | None, verbose: bool, cmd_input: bytes | None, mod_env: dict[str, str] | None, interactive: bool, log_prefix: str, ) -> Result: def __log(prio: int, *args): log(prio, log_prefix, *args) def __log_block(prio: int, title: str, block: bytes | str | None): if self.__caps & self.Caps.LogOutput: return if not block: return if isinstance(block, bytes): encoding = sys.stdout.encoding or 'utf-8' block = block.decode(encoding).strip() # Needed to pacify pyright: block can't be anything else at this point assert isinstance(block, str) delim = f'---- {title} ----' __log(prio, f',{delim}') for line in block.splitlines(): __log(prio, '|', line) __log(prio, f'`{delim}') if wd is not None and not self.__caps & self.Caps.Wd: cmd = ['cd', wd, '&&', *cmd] if interactive and not self.__caps & self.Caps.Interactive: raise NotImplementedError('Interactive SSH is not yet implemented') if mod_env is not None and not self.__caps & self.Caps.ModEnv: raise NotImplementedError( 'Passing an environment to SSH commands is not yet implemented' ) ret = await self._run_ssh( cmd = cmd, wd = wd, verbose = verbose, cmd_input = cmd_input, mod_env = mod_env, interactive = interactive, log_prefix = log_prefix, ) if verbose: __log_block(NOTICE, 'stdout', ret.stdout_str_or_none) __log_block(NOTICE, 'stderr', ret.stderr_str_or_none) if ret.status != 0: __log(WARNING, f'Exit code {ret.status}') return ret @property def hostname(self) -> str | None: return self.uri.hostname @property def port(self) -> int | None: return self.uri.port @property def username(self) -> str | None: return self.uri.username @property def password(self) -> str | None: return self.uri.password def ssh_client( *args, type: str | list[str] | None = None, **kwargs ) -> SSHClient: # export from importlib import import_module errors: list[str] = [] if type is None: val = os.getenv('JW_DEFAULT_SSH_CLIENT') if val is not None: type = val.split(',') else: type = ['AsyncSSH', 'Paramiko', 'Exec'] if isinstance(type, str): type = [type] for name in type: try: ret = getattr(import_module(f'jw.pkg.lib.ec.ssh.{name}'), name)(*args, **kwargs) log(INFO, f'Using SSH-client "{name}"') return ret except Exception as e: msg = f"Can't instantiate SSH client class {name} ({str(e)})" errors.append(msg) log(DEBUG, f'{msg}, trying next') msg = f'No working SSH clients for {" ".join([str(arg) for arg in args])}' log(ERR, f'----- {msg}') for error in errors: log(ERR, error) raise Exception(msg)