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

# -*- coding: utf-8 -*-

from typing import Any

import os, abc, shlex, sys

from ..util import run_cmd, pretty_cmd
from ..log import *
from ..ExecContext import ExecContext, Result
from urllib.parse import urlparse

class SSHClient(ExecContext):

    def __init__(self, uri: str, *args, **kwargs) -> None:
        super().__init__(uri=uri, *args, **kwargs)
        try:
            parsed = urlparse(uri)
        except Exception as e:
            log(ERR, f'Failed to parse SSH URI "{uri}"')
            raise
        self.__hostname = parsed.hostname
        self.__password: parsed.password
        self.__username = parsed.username

    @abc.abstractmethod
    async def _run_ssh(self, cmd: list[str]) -> Result:
        pass

    async def _run(
        self,
        cmd: list[str],
        wd: str|None,
        verbose: bool,
        cmd_input: str|None,
        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: str):
            encoding = sys.stdout.encoding or 'utf-8'
            block = block.decode(encoding).strip()
            if not block:
                return
            delim = f'---- {title} ----'
            __log(prio, f',{delim}')
            for line in block.splitlines():
                __log(prio, '|', line)
            __log(prio, f'`{delim}')

        if wd is not None:
            cmd = ['cd', wd, '&&', *cmd]

        if interactive:
            raise NotImplementedError('Interactive SSH is not yet implemented')

        if env is not None:
            raise NotImplementedError('Passing an environment to SSH commands is not yet implemented')

        ret = await self._run_ssh(cmd, cmd_input=cmd_input)
        if verbose:
            __log_block(NOTICE, 'stdout', ret.stdout)
            __log_block(NOTICE, 'stderr', ret.stderr)
            if ret.status != 0:
                __log(WARNING, f'Exit code {ret.status}')
        return ret

    async def _sudo(self, cmd: list[str], mod_env: dict[str, str], opts: list[str], *args, **kwargs) -> Result:
        if self.username != 'root':
            cmd = ['sudo', *opts, *cmd]
        if mod_env:
            log(WARNING, f'Modifying environment over SSH is not implemented, ignored')
        return await self._run(cmd, *args, **kwargs)

    @property
    def hostname(self):
        return self.__hostname

    def set_password(self, password: str) -> None:
        self.__password = password

    @property
    def password(self) -> str:
        return self.__password

    def set_username(self, username: str) -> None:
        self.__username = username

    @property
    def username(self) -> str:
        return self.__username

class SSHClientInternal(SSHClient): # export

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.__timeout: float|None = None # Untested
        self.___ssh: Any|None = None

    def __ssh_connect(self):
        import paramiko # type: ignore # error: Library stubs not installed for "paramiko"
        ret = paramiko.SSHClient()
        ret.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            ret.connect(
                hostname=self.hostname,
                username=self.username,
                allow_agent=True
            )
        except Exception as e:
            log(ERR, f'Failed to connect to {self.hostname} with key file {path_to_key} ({str(e)})')
            raise
        s = ret.get_transport().open_session()
        # set up the agent request handler to handle agent requests from the server
        paramiko.agent.AgentRequestHandler(s)
        return ret

    @property
    def __ssh(self):
        if self.___ssh is None:
            self.___ssh = self.__ssh_connect()
        return self.___ssh

    @property
    def __scp(self):
        return SCPClient(self.__ssh.get_transport())

    async def _run_ssh(self, cmd: list[str], cmd_input: str|None) -> Result:
        try:
            stdin, stdout, stderr = self.__ssh.exec_command(shlex.join(cmd), timeout=self.__timeout)
        except Exception as e:
            log(ERR, f'Command failed for {self.uri}: "{shlex.join(cmd)}"')
            raise
        if cmd_input is not None:
            stdin.write(cmd_input)
        exit_status = stdout.channel.recv_exit_status()
        return Result(stdout.read(), stderr.read(), exit_status)

class SSHClientCmd(SSHClient): # export

    def __init__(self, *args, **kwargs) -> None:
        self.__askpass: str|None = None
        self.__askpass_orig: dict[str, str|None] = dict()
        super().__init__(*args, **kwargs)

    def __del__(self):
        for key, val in self.__askpass_orig.items():
            if val is None:
                del os.environ[key]
            else:
                os.environ[key] = val
        if self.__askpass is not None:
            os.remove(self.__askpass)

    def __init_askpass(self):
        if self.__askpass is None and self.password is not None:
            import sys, tempfile
            prefix = os.path.basename(sys.argv[0]) + '-'
            f = tempfile.NamedTemporaryFile(mode='w+t', prefix=prefix, delete=False)
            os.chmod(f.name, 0o0700)
            self.__askpass = f.name
            f.write(f'#!/bin/bash\n\necho -n "{self.password}\n"')
            f.close()
            for key, val in {'SSH_ASKPASS': self.__askpass, 'SSH_ASKPASS_REQUIRE': 'force'}.items():
                self.__askpass_orig[key] = os.getenv(key)
                os.environ[key] = val

    async def _run_ssh(self, cmd: list[str], cmd_input: str|None) -> Result:
        self.__init_askpass()
        return await run_cmd(['ssh', self.hostname, shlex.join(cmd)],
                             output_encoding='bytes', cmd_input=cmd_input)

def ssh_client(*args, **kwargs) -> SSHClient: # export
    try:
        return SSHClientInternal(*args, **kwargs)
    except:
        pass
    return SSHClientCmd(*args, **kwargs)