lib.SSHClient: Move to lib.ec

SSHClient in an ExecContext, hence it's better off in lib.ec, move it
there and adapt the references.

Signed-off-by: Jan Lindemann <jan@janware.com>

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

@@ -0,0 +1,181 @@
+# -*- 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)