@neoline/hostbridge 1.4.4 → 2.0.0

package/src/server_control_mcp/transports/pty.py

@@ -0,0 +1,405 @@
+"""Ordered interactive PTY transport."""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import hashlib
+import re
+import shlex
+import threading
+import time
+import uuid
+from collections.abc import Callable, Iterable, Iterator
+from contextlib import suppress
+from typing import Protocol
+
+import pexpect
+
+from ..hosts import HostConfig
+from .base import ExecResult, ShellReadResult, TransferResult, TransportBusy, TransportCapabilities, TransportError
+
+_REMOTE_BASE64_DECODE = "{ base64 --decode 2>/dev/null || base64 -d 2>/dev/null || base64 -D; }"
+
+
+class ChildProcess(Protocol):
+    before: str
+    after: object
+
+    def sendline(self, line: str) -> None:
+        raise NotImplementedError
+
+    def send(self, text: str) -> None:
+        raise NotImplementedError
+
+    def expect(self, patterns: object, timeout: int | float = -1) -> int:
+        raise NotImplementedError
+
+    def isalive(self) -> bool:
+        raise NotImplementedError
+
+    def close(self, force: bool = False) -> None:
+        raise NotImplementedError
+
+    def read_nonblocking(self, size: int = 1, timeout: int | float = -1) -> str:
+        raise NotImplementedError
+
+
+ChildFactory = Callable[[str, list[str], int | float], ChildProcess]
+
+
+def _default_child_factory(command: str, args: list[str], timeout: int | float) -> ChildProcess:
+    return pexpect.spawn(
+        command,
+        args,
+        encoding="utf-8",
+        codec_errors="replace",
+        timeout=timeout,
+        echo=False,
+    )
+
+
+def _iter_base64(chunks: Iterable[bytes]) -> Iterator[str]:
+    remainder = b""
+    for chunk in chunks:
+        if not isinstance(chunk, bytes):
+            raise TypeError("stdin chunks must be bytes")
+        data = remainder + chunk
+        complete = len(data) - (len(data) % 3)
+        if complete:
+            yield base64.b64encode(data[:complete]).decode("ascii")
+        remainder = data[complete:]
+    if remainder:
+        yield base64.b64encode(remainder).decode("ascii")
+
+
+def _marker_text(value: object) -> str:
+    return value.group(0) if hasattr(value, "group") else str(value)
+
+
+class PtyTransport:
+    capabilities = TransportCapabilities(
+        separate_stderr=False,
+        native_binary_transfer=False,
+        parallel_channels=1,
+    )
+
+    @classmethod
+    def connect(
+        cls,
+        host: HostConfig,
+        *,
+        child_factory: ChildFactory = _default_child_factory,
+        connect_timeout: int | float = 60,
+        remote_shell: str = "bash",
+        token_factory: Callable[[], str] | None = None,
+    ) -> PtyTransport:
+        if connect_timeout <= 0:
+            raise ValueError("connect_timeout must be positive")
+        tokens = token_factory or (lambda: uuid.uuid4().hex)
+        child = child_factory(host.command, list(host.args), connect_timeout)
+        try:
+            for step in host.login_steps:
+                child.expect(step.expect, timeout=connect_timeout)
+                if step.send_secret is not None:
+                    secret = host.secrets.get(step.send_secret)
+                    if secret is None:
+                        raise TransportError(f"login step references missing secret {step.send_secret!r}")
+                    child.send(f"{secret}\r")
+                else:
+                    child.send(f"{step.send or ''}\r")
+            if host.login_steps:
+                child.expect(host.ready_expect or r"(?m)[^\r\n]*[#$]\s*$", timeout=connect_timeout)
+            token = tokens()
+            marker = f"__HB_READY_{token}__"
+            child.send(f"printf '\n{marker}\n'\r")
+            child.expect(re.escape(marker), timeout=connect_timeout)
+        except Exception as exc:
+            child.close(force=True)
+            if isinstance(exc, TransportError):
+                raise
+            raise TransportError(f"failed to establish PTY transport for {host.id}") from exc
+        return cls(child, remote_shell=remote_shell, token_factory=tokens)
+
+    def __init__(
+        self,
+        child: ChildProcess,
+        *,
+        remote_shell: str = "bash",
+        token_factory: Callable[[], str] | None = None,
+        marker_prefix: str = "HB",
+    ) -> None:
+        self.child = child
+        self.remote_shell = remote_shell
+        self._token_factory = token_factory or (lambda: uuid.uuid4().hex)
+        self._marker_prefix = marker_prefix
+        self.operation_lock = threading.Lock()
+        self._io_lock = threading.Lock()
+
+    def exec(
+        self,
+        command: str,
+        *,
+        stdin: Iterable[bytes] | None,
+        timeout: float,
+        output_limit: int,
+    ) -> ExecResult:
+        if not command.strip():
+            raise ValueError("command must not be empty")
+        if timeout <= 0:
+            raise ValueError("timeout must be positive")
+        if output_limit <= 0:
+            raise ValueError("output_limit must be positive")
+        if not self.operation_lock.acquire(blocking=False):
+            raise TransportBusy("PTY transport already has an active operation")
+        stdin_path: str | None = None
+        try:
+            if not self.child.isalive():
+                raise TransportError("PTY transport is not alive")
+            if stdin is not None:
+                stdin_path = self._stage_stdin_locked(stdin, timeout=timeout)
+                command = f"{command} < {shlex.quote(stdin_path)}"
+            return self._exec_locked(command, timeout=timeout, output_limit=output_limit)
+        finally:
+            if stdin_path is not None and self.child.isalive():
+                with suppress(Exception):
+                    self._exec_locked(
+                        f"rm -f -- {shlex.quote(stdin_path)}",
+                        timeout=min(timeout, 10),
+                        output_limit=1024,
+                    )
+            self.operation_lock.release()
+
+    def _stage_stdin_locked(self, chunks: Iterable[bytes], *, timeout: float) -> str:
+        token = self._token_factory()
+        path = f"/tmp/hostbridge.stdin.{token}"
+        heredoc = f"__{self._marker_prefix}_STDIN_{token}__"
+        done_pattern = re.compile(rf"__{self._marker_prefix}_STDIN_DONE_{re.escape(token)}__:(\d+)")
+        self.child.send(f"cat > {shlex.quote(path)} <<'{heredoc}'\n")
+        for encoded in _iter_base64(chunks):
+            self.child.send(encoded)
+        self.child.send(
+            f"\n{heredoc}\n"
+            f"tmp={shlex.quote(path)}.b64; mv {shlex.quote(path)} \"$tmp\"; "
+            f"{{ {_REMOTE_BASE64_DECODE} < \"$tmp\"; }} > {shlex.quote(path)}; status=$?; rm -f \"$tmp\"; "
+            f"printf '\n__{self._marker_prefix}_STDIN_DONE_{token}__:%s\n' \"$status\"\n"
+        )
+        self.child.expect(done_pattern, timeout=timeout)
+        match = re.search(r":(\d+)", _marker_text(self.child.after))
+        if match is None or int(match.group(1)) != 0:
+            raise TransportError("failed to stage PTY stdin")
+        return path
+
+    def _exec_locked(self, command: str, *, timeout: float, output_limit: int) -> ExecResult:
+        token = self._token_factory()
+        start_marker = f"__{self._marker_prefix}_START_{token}__"
+        exit_pattern = re.compile(rf"__{self._marker_prefix}_EXIT_{re.escape(token)}__:(\d+)")
+        encoded_command = base64.b64encode(command.encode("utf-8")).decode("ascii")
+        remote_line = (
+            f"__scm_cmd=$(printf %s {shlex.quote(encoded_command)} | {_REMOTE_BASE64_DECODE}); "
+            f"printf '\n{start_marker}\n'; "
+            f"{self.remote_shell} -lc \"$__scm_cmd\"; "
+            "__scm_status=$?; "
+            f"printf '\n__{self._marker_prefix}_EXIT_{token}__:%s\n' \"$__scm_status\""
+        )
+        try:
+            self.child.sendline(remote_line)
+            self.child.expect(exit_pattern, timeout=timeout)
+        except pexpect.TIMEOUT:
+            return ExecResult(None, self._output_after_marker(str(self.child.before), start_marker, output_limit), timed_out=True)
+        match = re.search(r":(\d+)", _marker_text(self.child.after))
+        exit_code = int(match.group(1)) if match else None
+        return ExecResult(exit_code, self._output_after_marker(str(self.child.before), start_marker, output_limit))
+
+    @staticmethod
+    def _output_after_marker(raw: str, marker: str, output_limit: int) -> bytes:
+        normalized = raw.replace("\x1b[?2004h", "").replace("\x1b[?2004l", "")
+        normalized = normalized.replace("\r\n", "\n").replace("\r", "\n")
+        marker_index = normalized.rfind(marker)
+        output = normalized[marker_index + len(marker) :] if marker_index >= 0 else normalized
+        cleaned = [
+            line
+            for line in output.split("\n")
+            if not line.startswith("printf '") and not re.search(r"__[A-Z]+_(?:START|EXIT)_", line)
+        ]
+        encoded = "\n".join(cleaned).strip("\n").encode("utf-8", errors="replace")
+        return encoded[-output_limit:]
+
+    def upload(
+        self,
+        chunks: Iterable[bytes],
+        remote_path: str,
+        *,
+        mode: int,
+        expected_size: int | None,
+    ) -> TransferResult:
+        self._validate_remote_path(remote_path)
+        if not isinstance(mode, int) or isinstance(mode, bool) or not 0 <= mode <= 0o7777:
+            raise ValueError("mode must be an integer between 0 and 07777")
+        if expected_size is not None and expected_size < 0:
+            raise ValueError("expected_size must be non-negative")
+        if not self.operation_lock.acquire(blocking=False):
+            raise TransportBusy("PTY transport already has an active operation")
+        token = self._token_factory()
+        tmp_b64 = f"/tmp/hostbridge.upload.{token}.b64"
+        tmp_out = f"/tmp/hostbridge.upload.{token}.out"
+        heredoc = f"__{self._marker_prefix}_UPLOAD_{token}__"
+        done_pattern = re.compile(rf"__{self._marker_prefix}_UPLOAD_DONE_{re.escape(token)}__:(\d+)")
+        digest = hashlib.sha256()
+        transferred = 0
+
+        def counted_chunks() -> Iterator[bytes]:
+            nonlocal transferred
+            for chunk in chunks:
+                if not isinstance(chunk, bytes):
+                    raise TypeError("upload chunks must be bytes")
+                transferred += len(chunk)
+                digest.update(chunk)
+                yield chunk
+
+        try:
+            if not self.child.isalive():
+                raise TransportError("PTY transport is not alive")
+            self.child.send(f"cat > {shlex.quote(tmp_b64)} <<'{heredoc}'\n")
+            for encoded in _iter_base64(counted_chunks()):
+                self.child.send(encoded)
+            self.child.send(
+                f"\n{heredoc}\nprintf '\n__{self._marker_prefix}_UPLOAD_DONE_{token}__:%s\n' \"$?\"\n"
+            )
+            self.child.expect(done_pattern, timeout=60)
+            match = re.search(r":(\d+)", _marker_text(self.child.after))
+            if match is None or int(match.group(1)) != 0:
+                raise TransportError("failed to stream PTY upload")
+            if expected_size is not None and transferred != expected_size:
+                raise TransportError(
+                    f"upload size mismatch: expected {expected_size}, received {transferred}"
+                )
+            target_q = shlex.quote(remote_path)
+            finalize = (
+                "set -e; "
+                f"target={target_q}; tmp_b64={shlex.quote(tmp_b64)}; tmp_out={shlex.quote(tmp_out)}; "
+                'mkdir -p "$(dirname "$target")"; '
+                f'{{ {_REMOTE_BASE64_DECODE} < "$tmp_b64"; }} > "$tmp_out"; '
+                f'chmod {format(mode, "04o")} "$tmp_out"; mv "$tmp_out" "$target"; rm -f "$tmp_b64"; '
+                "if command -v sha256sum >/dev/null 2>&1; then sha=$(sha256sum \"$target\" | awk '{print $1}'); "
+                "else sha=$(shasum -a 256 \"$target\" | awk '{print $1}'); fi; "
+                'bytes=$(wc -c < "$target" | tr -d " "); printf "BYTES=%s\\nSHA256=%s\\n" "$bytes" "$sha"'
+            )
+            result = self._exec_locked(finalize, timeout=60, output_limit=4096)
+            if result.timed_out or result.exit_code != 0:
+                raise TransportError("failed to finalize PTY upload")
+            fields = self._metadata_fields(result.stdout)
+            remote_size = int(fields.get("BYTES", transferred))
+            remote_sha = fields.get("SHA256", digest.hexdigest())
+            if remote_size != transferred or remote_sha != digest.hexdigest():
+                raise TransportError("PTY upload integrity check failed")
+            return TransferResult(transferred, remote_sha)
+        except Exception:
+            if self.child.isalive():
+                with suppress(Exception):
+                    self._exec_locked(
+                        f"rm -f -- {shlex.quote(tmp_b64)} {shlex.quote(tmp_out)}",
+                        timeout=10,
+                        output_limit=1024,
+                    )
+            raise
+        finally:
+            self.operation_lock.release()
+
+    def download(self, remote_path: str, *, chunk_size: int) -> Iterator[bytes]:
+        self._validate_remote_path(remote_path)
+        if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= 1024 * 1024:
+            raise ValueError("chunk_size must be between 4096 and 1048576 bytes")
+
+        def stream() -> Iterator[bytes]:
+            if not self.operation_lock.acquire(blocking=False):
+                raise TransportBusy("PTY transport already has an active operation")
+            digest = hashlib.sha256()
+            transferred = 0
+            try:
+                if not self.child.isalive():
+                    raise TransportError("PTY transport is not alive")
+                remote_q = shlex.quote(remote_path)
+                metadata = (
+                    f"target={remote_q}; bytes=$(wc -c < \"$target\" | tr -d ' '); "
+                    "if command -v sha256sum >/dev/null 2>&1; then sha=$(sha256sum \"$target\" | awk '{print $1}'); "
+                    "else sha=$(shasum -a 256 \"$target\" | awk '{print $1}'); fi; "
+                    'printf "BYTES=%s\\nSHA256=%s\\n" "$bytes" "$sha"'
+                )
+                result = self._exec_locked(metadata, timeout=60, output_limit=4096)
+                if result.timed_out or result.exit_code != 0:
+                    raise TransportError(f"failed to inspect remote file {remote_path}")
+                fields = self._metadata_fields(result.stdout)
+                byte_count = int(fields.get("BYTES", "0"))
+                expected_sha = fields.get("SHA256", "")
+                chunk_count = (byte_count + chunk_size - 1) // chunk_size
+                for index in range(chunk_count):
+                    command = (
+                        f"dd if={remote_q} bs={chunk_size} skip={index} count=1 2>/dev/null | "
+                        "{ base64 --wrap=0 2>/dev/null || base64 | tr -d '\\n'; }"
+                    )
+                    result = self._exec_locked(command, timeout=60, output_limit=chunk_size * 2 + 4096)
+                    if result.timed_out or result.exit_code != 0:
+                        raise TransportError(f"failed to download PTY chunk {index}")
+                    try:
+                        chunk = base64.b64decode(b"".join(result.stdout.split()), validate=True)
+                    except (binascii.Error, ValueError) as exc:
+                        raise TransportError(f"remote chunk {index} contains invalid base64") from exc
+                    transferred += len(chunk)
+                    digest.update(chunk)
+                    yield chunk
+                if transferred != byte_count or (expected_sha and digest.hexdigest() != expected_sha):
+                    raise TransportError("PTY download integrity check failed")
+            finally:
+                self.operation_lock.release()
+
+        return stream()
+
+    @staticmethod
+    def _validate_remote_path(remote_path: str) -> None:
+        if not isinstance(remote_path, str) or not remote_path.strip() or "\x00" in remote_path:
+            raise ValueError("remote_path must be a non-empty path without NUL bytes")
+
+    @staticmethod
+    def _metadata_fields(output: bytes) -> dict[str, str]:
+        fields: dict[str, str] = {}
+        for line in output.decode("utf-8", errors="replace").splitlines():
+            key, separator, value = line.partition("=")
+            if separator:
+                fields[key] = value
+        return fields
+
+    def close(self) -> None:
+        self.child.close(force=True)
+
+    def shell_write(self, data: bytes) -> None:
+        if not isinstance(data, bytes):
+            raise TypeError("shell data must be bytes")
+        if not self.child.isalive():
+            raise TransportError("PTY transport is not alive")
+        with self._io_lock:
+            self.child.send(data.decode("utf-8", errors="replace"))
+
+    def shell_read(self, *, timeout: float, max_bytes: int) -> ShellReadResult:
+        if timeout < 0:
+            raise ValueError("timeout must be non-negative")
+        if max_bytes <= 0:
+            raise ValueError("max_bytes must be positive")
+        output = bytearray()
+        deadline = time.monotonic() + timeout
+        with self._io_lock:
+            while len(output) < max_bytes:
+                remaining = max(0.0, deadline - time.monotonic())
+                try:
+                    chunk = self.child.read_nonblocking(
+                        size=min(4096, max_bytes - len(output)),
+                        timeout=remaining,
+                    )
+                except (pexpect.TIMEOUT, pexpect.EOF):
+                    break
+                if not chunk:
+                    break
+                output.extend(chunk.encode("utf-8", errors="replace"))
+                if remaining <= 0:
+                    break
+        return ShellReadResult(bytes(output), self.child.isalive())

package/src/server_control_mcp/transports/ssh.py

@@ -0,0 +1,195 @@
+"""Native AsyncSSH transport for direct SSH hosts."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import uuid
+from collections.abc import Iterable, Sequence
+from contextlib import suppress
+from dataclasses import dataclass
+from pathlib import Path
+
+import asyncssh
+
+from .base import ExecResult, TransferResult, TransportCapabilities
+
+
+@dataclass(frozen=True, slots=True)
+class SshConnectionConfig:
+    host: str
+    port: int = 22
+    username: str | None = None
+    client_keys: Sequence[str] = ()
+    known_hosts: str | None = None
+
+
+class SshTransport:
+    capabilities = TransportCapabilities(
+        separate_stderr=True,
+        native_binary_transfer=True,
+        parallel_channels=8,
+    )
+
+    def __init__(self, connection: asyncssh.SSHClientConnection) -> None:
+        self._connection = connection
+
+    @classmethod
+    async def connect(cls, config: SshConnectionConfig) -> SshTransport:
+        if not config.host.strip():
+            raise ValueError("SSH host must not be empty")
+        if not 1 <= config.port <= 65535:
+            raise ValueError("SSH port must be between 1 and 65535")
+        connection = await asyncssh.connect(
+            config.host,
+            port=config.port,
+            username=config.username,
+            client_keys=list(config.client_keys),
+            known_hosts=config.known_hosts,
+            encoding=None,
+        )
+        return cls(connection)
+
+    async def exec(
+        self,
+        command: str,
+        *,
+        stdin: Iterable[bytes] | None,
+        timeout: float,
+        output_limit: int,
+    ) -> ExecResult:
+        if not command.strip():
+            raise ValueError("command must not be empty")
+        if timeout <= 0:
+            raise ValueError("timeout must be positive")
+        if output_limit <= 0:
+            raise ValueError("output_limit must be positive")
+        process = await self._connection.create_process(command, encoding=None)
+        stdout_task = asyncio.create_task(self._read_stream(process.stdout, output_limit))
+        stderr_task = asyncio.create_task(self._read_stream(process.stderr, output_limit))
+        input_task = asyncio.create_task(self._write_stdin(process, stdin))
+        try:
+            await asyncio.wait_for(process.wait_closed(), timeout=timeout)
+            await input_task
+            stdout, stderr = await asyncio.gather(stdout_task, stderr_task)
+            return ExecResult(process.exit_status, stdout, stderr)
+        except TimeoutError:
+            process.terminate()
+            with suppress(Exception):
+                await process.wait_closed()
+            for task in (stdout_task, stderr_task, input_task):
+                task.cancel()
+            await asyncio.gather(stdout_task, stderr_task, input_task, return_exceptions=True)
+            return ExecResult(None, b"", b"", timed_out=True)
+
+    @staticmethod
+    async def _write_stdin(process: asyncssh.SSHClientProcess[bytes], stdin: Iterable[bytes] | None) -> None:
+        if stdin is not None:
+            for chunk in stdin:
+                if not isinstance(chunk, bytes):
+                    raise TypeError("stdin chunks must be bytes")
+                process.stdin.write(chunk)
+                await process.stdin.drain()
+        process.stdin.write_eof()
+
+    @staticmethod
+    async def _read_stream(stream, output_limit: int) -> bytes:
+        output = bytearray()
+        while True:
+            chunk = await stream.read(64 * 1024)
+            if not chunk:
+                break
+            output.extend(chunk)
+            if len(output) > output_limit:
+                del output[: len(output) - output_limit]
+        return bytes(output)
+
+    async def upload_file(
+        self,
+        local_path: Path | str,
+        remote_path: str,
+        *,
+        chunk_size: int = 256 * 1024,
+        mode: int = 0o644,
+    ) -> TransferResult:
+        source = Path(local_path)
+        self._validate_chunk_size(chunk_size)
+        self._validate_remote_path(remote_path)
+        temporary = f"{remote_path}.hostbridge-{uuid.uuid4().hex}.tmp"
+        digest = hashlib.sha256()
+        transferred = 0
+        sftp = await self._connection.start_sftp_client()
+        try:
+            local_file = source.open("rb")
+            try:
+                async with sftp.open(temporary, "wb") as remote_file:
+                    while True:
+                        chunk = await asyncio.to_thread(local_file.read, chunk_size)
+                        if not chunk:
+                            break
+                        await remote_file.write(chunk)
+                        digest.update(chunk)
+                        transferred += len(chunk)
+            finally:
+                local_file.close()
+            await sftp.chmod(temporary, mode)
+            await sftp.rename(temporary, remote_path)
+        except Exception:
+            with suppress(Exception):
+                await sftp.remove(temporary)
+            raise
+        finally:
+            sftp.exit()
+        return TransferResult(transferred, digest.hexdigest())
+
+    async def download_file(
+        self,
+        remote_path: str,
+        local_path: Path | str,
+        *,
+        chunk_size: int = 256 * 1024,
+    ) -> TransferResult:
+        destination = Path(local_path)
+        self._validate_chunk_size(chunk_size)
+        self._validate_remote_path(remote_path)
+        temporary = destination.with_name(f".{destination.name}.hostbridge-{uuid.uuid4().hex}.tmp")
+        digest = hashlib.sha256()
+        transferred = 0
+        sftp = await self._connection.start_sftp_client()
+        try:
+            destination.parent.mkdir(parents=True, exist_ok=True)
+            local_file = temporary.open("wb")
+            try:
+                async with sftp.open(remote_path, "rb") as remote_file:
+                    while True:
+                        raw_chunk = await remote_file.read(chunk_size)
+                        if not raw_chunk:
+                            break
+                        chunk = raw_chunk.encode("utf-8") if isinstance(raw_chunk, str) else bytes(raw_chunk)
+                        await asyncio.to_thread(local_file.write, chunk)
+                        digest.update(chunk)
+                        transferred += len(chunk)
+            finally:
+                local_file.close()
+            temporary.chmod(0o600)
+            temporary.replace(destination)
+        except Exception:
+            temporary.unlink(missing_ok=True)
+            raise
+        finally:
+            sftp.exit()
+        return TransferResult(transferred, digest.hexdigest())
+
+    @staticmethod
+    def _validate_chunk_size(chunk_size: int) -> None:
+        if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= 1024 * 1024:
+            raise ValueError("chunk_size must be between 4096 and 1048576 bytes")
+
+    @staticmethod
+    def _validate_remote_path(remote_path: str) -> None:
+        if not isinstance(remote_path, str) or not remote_path.strip() or "\x00" in remote_path:
+            raise ValueError("remote_path must be a non-empty path without NUL bytes")
+
+    async def close(self) -> None:
+        self._connection.close()
+        await self._connection.wait_closed()