from __future__ import annotations import abc import os import pwd from enum import Flag, auto from typing import TYPE_CHECKING, Any, cast, override from ..ExecContext import ExecContext from ..log import DEBUG, ERR, INFO, NOTICE, get_caller_pos, log, log_m from ..Uri import Uri if TYPE_CHECKING: from ..base import Result 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: Any, **kwargs: Any ) -> 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 @override 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: Any, **kwargs: Any) -> None: caller = kwargs.get('caller') if caller is None: kwargs['caller'] = get_caller_pos(1) log(prio, log_prefix, *args, **kwargs) def __log_block(prio: int, title: str, block: str | None) -> None: if self.__caps & self.Caps.LogOutput: return if block is None: return log_m(prio, f'---- {title} ----\n{block}', caller = get_caller_pos(1)) 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) return ret @property def hostname(self) -> str | None: return self.uri.hostname @property def port(self) -> int | None: return self.uri.port @property @override def username(self) -> str | None: return self.uri.username @property def password(self) -> str | None: return self.uri.password def ssh_client( # export *args: Any, type: str | list[str] | None = None, **kwargs: Any ) -> 'SSHClient': 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 cast('SSHClient', 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)