jw-pkg/src/python/jw/pkg/lib/ec/ssh/AsyncSSH.py

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

import os, sys, shlex, asyncio, asyncssh, shutil, signal

from ...log import *
from ...base import Result
from ..SSHClient import SSHClient as Base

from .util import join_cmd

_USE_DEFAULT_KNOWN_HOSTS = object()

class AsyncSSH(Base):

    def __init__(
        self,
        uri: str,
        *,
        client_keys: list[str] | None = None,
        known_hosts = _USE_DEFAULT_KNOWN_HOSTS,
        term_type: str | None = None,
        connect_timeout: float | None = 30.0,
        **kwargs,
    ) -> None:

        super().__init__(
            uri,
            caps = self.Caps.LogOutput | self.Caps.Wd | self.Caps.Interactive | self.Caps.Env,
            **kwargs
        )

        self.client_keys = client_keys
        self.known_hosts = known_hosts
        self.term_type = term_type or os.environ.get('TERM', 'xterm')
        self.connect_timeout = connect_timeout

    def _connect_kwargs(self, hide_secrets: bool=False) -> dict:

        kwargs: dict = {
            'host': self.hostname,
            'port': self.port,
            'username': self.username,
            'password': self.password,
            'client_keys': self.client_keys,
            'connect_timeout': self.connect_timeout,
        }

        if self.known_hosts is not _USE_DEFAULT_KNOWN_HOSTS:
            kwargs['known_hosts'] = self.known_hosts

        ret = {k: v for k, v in kwargs.items() if v is not None}
        if hide_secrets and 'password' in kwargs:
            kwargs['password'] = '<hidden>'
        return ret

    @staticmethod
    def _build_remote_command(cmd: list[str], wd: str | None) -> str:

        if not cmd:
            raise ValueError('cmd must not be empty')

        inner = f'exec {join_cmd(cmd)}'
        if wd is not None:
            inner = f'cd {shlex.quote(wd)} && {inner}'

        return f'/bin/sh -lc {shlex.quote(inner)}'

    @staticmethod
    def _has_local_tty() -> bool:
        try:
            return sys.stdin.isatty() and sys.stdout.isatty()
        except Exception:
            return False

    @staticmethod
    def _get_local_term_size() -> tuple[int, int, int, int]:
        cols, rows = shutil.get_terminal_size(fallback=(80, 24))
        xpixel = ypixel = 0

        try:
            import fcntl, termios, struct

            packed = fcntl.ioctl(sys.stdout.fileno(), termios.TIOCGWINSZ, b'\0' * 8)
            rows2, cols2, xpixel, ypixel = struct.unpack('HHHH', packed)

            if cols2 > 0 and rows2 > 0:
                cols, rows = cols2, rows2
        except Exception:
            pass

        return (cols, rows, xpixel, ypixel)

    async def _read_stream(
        self,
        stream,
        prio,
        collector: list[bytes],
        *,
        verbose: bool,
        log_prefix: str,
        log_enc: str,
    ) -> None:
        buf = b''

        while True:
            chunk = await stream.read(4096)
            if not chunk:
                break

            collector.append(chunk)

            if verbose:
                buf += chunk
                while b'\n' in buf:
                    line, buf = buf.split(b'\n', 1)
                    log(prio, log_prefix, line.decode(log_enc, errors='replace'))

        if verbose and buf:
            log(prio, log_prefix, buf.decode(log_enc, errors='replace'))

    async def _run_interactive_on_conn(
        self,
        conn: asyncssh.SSHClientConnection,
        cmd: list[str],
        wd: str | None,
        cmd_input: bytes | None,
        mod_env: dict[str, str] | None,
    ) -> Result:

        command = self._build_remote_command(cmd, wd)
        stdout_parts: list[bytes] = []

        proc = await conn.create_process(
            command = command,
            env = mod_env,
            stdin = asyncssh.PIPE,
            stdout = asyncssh.PIPE,
            stderr = asyncssh.STDOUT,
            encoding = None,
            request_pty = 'force',
            term_type = self.term_type,
            term_size = self._get_local_term_size(),
        )

        loop = asyncio.get_running_loop()
        stdin_fd = sys.stdin.fileno()

        stdin_queue: asyncio.Queue[bytes | None] = asyncio.Queue()
        old_tty_state = None
        old_winch_handler = None
        stdin_reader_installed = False

        def _write_local(data: bytes) -> None:
            try:
                sys.stdout.buffer.write(data)
                sys.stdout.buffer.flush()
            except AttributeError:
                os.write(sys.stdout.fileno(), data)

        def _on_stdin_ready() -> None:
            try:
                data = os.read(stdin_fd, 4096)
            except OSError:
                data = b''

            if data:
                stdin_queue.put_nowait(data)
            else:
                try:
                    loop.remove_reader(stdin_fd)
                except Exception:
                    pass
                stdin_queue.put_nowait(None)

        async def _pump_stdin() -> None:
            if cmd_input is not None and proc.stdin is not None:
                proc.stdin.write(cmd_input)
                await proc.stdin.drain()
            while True:
                data = await stdin_queue.get()
                if data is None:
                    if proc.stdin is not None:
                        try:
                            proc.stdin.write_eof()
                        except (BrokenPipeError, OSError):
                            pass
                    return
                if proc.stdin is None:
                    return
                proc.stdin.write(data)
                await proc.stdin.drain()

        async def _pump_stdout() -> None:
            while True:
                chunk = await proc.stdout.read(4096)
                if not chunk:
                    break
                stdout_parts.append(chunk)
                _write_local(chunk)

        def _on_winch(*_args) -> None:

            try:
                proc.change_terminal_size(*self._get_local_term_size())
            except Exception:
                pass

        try:
            sys.stdout.flush()
            sys.stderr.flush()

            try:
                import termios, tty

                old_tty_state = termios.tcgetattr(stdin_fd)
                tty.setraw(stdin_fd)
            except Exception:
                old_tty_state = None

            try:
                loop.add_reader(stdin_fd, _on_stdin_ready)
                stdin_reader_installed = True
            except (NotImplementedError, RuntimeError):
                stdin_queue.put_nowait(None)

            if hasattr(signal, 'SIGWINCH'):
                try:
                    old_winch_handler = signal.getsignal(signal.SIGWINCH)
                    signal.signal(signal.SIGWINCH, _on_winch)
                except Exception:
                    old_winch_handler = None

            stdin_task = asyncio.create_task(_pump_stdin())
            stdout_task = asyncio.create_task(_pump_stdout())

            completed = await proc.wait(check = False)
            await stdout_task

            if not stdin_task.done():
                stdin_task.cancel()
                try:
                    await stdin_task
                except asyncio.CancelledError:
                    pass

            exit_code = completed.exit_status
            if exit_code is None:
                exit_code = completed.returncode if completed.returncode is not None else -1

            stdout = b''.join(stdout_parts) if stdout_parts else None
            return Result(stdout, None, exit_code)

        finally:
            if stdin_reader_installed:
                try:
                    loop.remove_reader(stdin_fd)
                except Exception:
                    pass

            if old_winch_handler is not None and hasattr(signal, 'SIGWINCH'):
                try:
                    signal.signal(signal.SIGWINCH, old_winch_handler)
                except Exception:
                    pass

            if old_tty_state is not None:
                try:
                    import termios
                    termios.tcsetattr(stdin_fd, termios.TCSADRAIN, old_tty_state)
                except Exception:
                    pass

            try:
                sys.stdout.flush()
                sys.stderr.flush()
            except Exception:
                pass

    async def _run_captured_pty_on_conn(
        self,
        conn: asyncssh.SSHClientConnection,
        cmd: list[str],
        wd: str | None,
        verbose: bool,
        cmd_input: bytes | None,
        mod_env: dict[str, str] | None,
        log_prefix: str,
    ) -> Result:
        command = self._build_remote_command(cmd, wd)

        stdout_parts: list[bytes] = []
        stdout_log_enc = sys.stdout.encoding or 'utf-8'

        proc = await conn.create_process(
            command = command,
            env = mod_env,
            stdin = asyncssh.PIPE if cmd_input is not None else asyncssh.DEVNULL,
            stdout = asyncssh.PIPE,
            stderr = asyncssh.STDOUT,
            encoding = None,
            request_pty = 'force',
            term_type = self.term_type,
        )

        task = asyncio.create_task(
            self._read_stream(
                proc.stdout,
                NOTICE,
                stdout_parts,
                verbose = verbose,
                log_prefix = log_prefix,
                log_enc = stdout_log_enc,
            )
        )

        if cmd_input is not None and proc.stdin is not None:
            proc.stdin.write(cmd_input)
            await proc.stdin.drain()
            proc.stdin.write_eof()

        completed = await proc.wait(check=False)
        await task

        exit_code = completed.exit_status
        if exit_code is None:
            exit_code = completed.returncode if completed.returncode is not None else -1

        stdout = b''.join(stdout_parts) if stdout_parts else None
        return Result(stdout, None, exit_code)

    async def _run_on_conn(
        self,
        conn: asyncssh.SSHClientConnection,
        cmd: list[str],
        wd: str | None,
        verbose: bool,
        cmd_input: bytes | None,
        mod_env: dict[str, str] | None,
        interactive: bool,
        log_prefix: str,
    ) -> Result:
        if interactive:
            if self._has_local_tty():
                return await self._run_interactive_on_conn(
                    conn = conn,
                    cmd = cmd,
                    wd = wd,
                    cmd_input = cmd_input,
                    mod_env = mod_env,
                )

            return await self._run_captured_pty_on_conn(
                conn = conn,
                cmd = cmd,
                wd = wd,
                verbose = verbose,
                cmd_input = cmd_input,
                mod_env = mod_env,
                log_prefix = log_prefix,
            )

        command = self._build_remote_command(cmd, wd)

        stdout_parts: list[bytes] = []
        stderr_parts: list[bytes] = []

        stdout_log_enc = sys.stdout.encoding or 'utf-8'
        stderr_log_enc = sys.stderr.encoding or 'utf-8'

        stdin_mode = asyncssh.PIPE if cmd_input is not None else asyncssh.DEVNULL

        proc = await conn.create_process(
            command = command,
            env = mod_env,
            stdin = stdin_mode,
            stdout = asyncssh.PIPE,
            stderr = asyncssh.PIPE,
            encoding = None,
            request_pty = False,
        )

        tasks = [
            asyncio.create_task(
                self._read_stream(
                    proc.stdout,
                    NOTICE,
                    stdout_parts,
                    verbose = verbose,
                    log_prefix = log_prefix,
                    log_enc = stdout_log_enc,
                )
            ),
            asyncio.create_task(
                self._read_stream(
                    proc.stderr,
                    ERR,
                    stderr_parts,
                    verbose = verbose,
                    log_prefix = log_prefix,
                    log_enc = stderr_log_enc,
                )
            ),
        ]

        if cmd_input is not None and proc.stdin is not None:
            proc.stdin.write(cmd_input)
            await proc.stdin.drain()
            proc.stdin.write_eof()

        completed = await proc.wait(check=False)
        await asyncio.gather(*tasks)

        stdout = b''.join(stdout_parts) if stdout_parts else None
        stderr = b''.join(stderr_parts) if stderr_parts else None

        exit_code = completed.exit_status
        if exit_code is None:
            exit_code = completed.returncode if completed.returncode is not None else -1

        return Result(stdout, stderr, exit_code)

    async def _run_ssh(
        self,
        cmd: list[str],
        wd: str | None,
        verbose: bool,
        cmd_input: str | None,
        mod_env: dict[str, str] | None,
        interactive: bool,
        log_prefix: str,
    ) -> Result:
        try:
            async with asyncssh.connect(**self._connect_kwargs()) as conn:
                return await self._run_on_conn(
                    conn = conn,
                    cmd = cmd,
                    wd = wd,
                    verbose = verbose,
                    cmd_input = cmd_input,
                    mod_env = mod_env,
                    interactive = interactive,
                    log_prefix = log_prefix,
                )
        except Exception as e:
            msg = f'-------------------- Failed to run command {" ".join(cmd)} ({e})'
            log(ERR, ',', msg)
            for key, val in self._connect_kwargs(hide_secrets=True).items():
                log(ERR, f'| {key:<20} = {val}')
            log(ERR, '`', msg)
            raise