jw-pkg/src/python/jw/pkg/lib/ec/SSHClient.py

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)