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) -> 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

        return {k: v for k, v in kwargs.items() if v is not None}

    @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 _merge_env_into_forwarded_args(
        args: tuple,
        kwargs: dict,
        mod_env: dict[str, str],
    ) -> tuple[tuple, dict]:
        args = list(args)
        kwargs = dict(kwargs)

        if "env" in kwargs:
            base_env = kwargs["env"]
            merged_env = dict(base_env or {})
            merged_env.update(mod_env)
            kwargs["env"] = merged_env or None
        elif len(args) >= 4:
            base_env = args[3]
            merged_env = dict(base_env or {})
            merged_env.update(mod_env)
            args[3] = merged_env or None
        else:
            kwargs["env"] = dict(mod_env) if mod_env else None
        return tuple(args), kwargs

    @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,
        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=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,
        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=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,
        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,
                    env=env,
                )

            return await self._run_captured_pty_on_conn(
                conn=conn,
                cmd=cmd,
                wd=wd,
                verbose=verbose,
                cmd_input=cmd_input,
                env=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=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,
        env: dict[str, str] | None,
        interactive: bool,
        log_prefix: str,
    ) -> Result:
        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,
                env=env,
                interactive=interactive,
                log_prefix=log_prefix,
            )

    async def _sudo(
        self,
        cmd: list[str],
        mod_env: dict[str, str],
        opts: list[str],
        *args,
        **kwargs,
    ) -> Result:
        args, kwargs = self._merge_env_into_forwarded_args(args, kwargs, mod_env)

        async with asyncssh.connect(**self._connect_kwargs()) as conn:
            uid_result = await conn.run("id -u", check=False)
            is_root = (
                uid_result.exit_status == 0
                and isinstance(uid_result.stdout, str)
                and uid_result.stdout.strip() == "0"
            )

            cmdline: list[str] = []

            if not is_root:
                cmdline.append("/usr/bin/sudo")
                if mod_env:
                    cmdline.append("--preserve-env=" + ",".join(mod_env.keys()))
                cmdline.extend(opts)

            cmdline.extend(cmd)

            return await self._run_on_conn(conn, cmdline, *args, **kwargs)