@neoline/hostbridge 2.0.3 → 2.0.5

package/src/server_control_mcp/transports/pty.py

@@ -1,434 +0,0 @@
-"""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; }"
-PTY_BASE64_LINE_LENGTH = 3072
-
-
-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]:
-    def split_line(encoded: str) -> Iterator[str]:
-        for offset in range(0, len(encoded), PTY_BASE64_LINE_LENGTH):
-            yield encoded[offset : offset + PTY_BASE64_LINE_LENGTH]
-
-    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 from split_line(base64.b64encode(data[:complete]).decode("ascii"))
-        remainder = data[complete:]
-    if remainder:
-        yield from split_line(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)
-            return self._exec_locked(command, timeout=timeout, output_limit=output_limit, stdin_path=stdin_path)
-        finally:
-            if stdin_path is not None and self.child.isalive():
-                with suppress(Exception):
-                    self._exec_locked(
-                        f"trap - EXIT HUP INT TERM; rm -f -- {shlex.quote(stdin_path)} {shlex.quote(f'{stdin_path}.b64')}",
-                        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}"
-        base64_path = f"{path}.b64"
-        heredoc = f"__{self._marker_prefix}_STDIN_{token}__"
-        done_pattern = re.compile(rf"__{self._marker_prefix}_STDIN_DONE_{re.escape(token)}__:(\d+)")
-        cleanup = f"rm -f -- {shlex.quote(path)} {shlex.quote(base64_path)}"
-        try:
-            self.child.send(f"trap {shlex.quote(cleanup)} EXIT HUP INT TERM; 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(base64_path)}; 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)
-        except BaseException:
-            if self.child.isalive():
-                with suppress(Exception):
-                    self.child.send("\x03")
-                with suppress(Exception):
-                    self.child.sendline(f"trap - EXIT HUP INT TERM; {cleanup}")
-            raise
-        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,
-        stdin_path: str | None = None,
-    ) -> 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")
-        stdin_redirect = f" < {shlex.quote(stdin_path)}" if stdin_path is not None else ""
-        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\"{stdin_redirect}; "
-            "__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}; "
-                    'if [ ! -e "$target" ] && [ ! -L "$target" ]; then printf "__HB_FILE_NOT_FOUND__\\n"; exit 44; fi; '
-                    'if [ ! -r "$target" ]; then printf "__HB_PERMISSION_DENIED__\\n"; exit 45; fi; '
-                    "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:
-                    if b"__HB_FILE_NOT_FOUND__" in result.stdout:
-                        raise TransportError(f"remote file not found: {remote_path}")
-                    if b"__HB_PERMISSION_DENIED__" in result.stdout:
-                        raise TransportError(f"permission denied reading remote file: {remote_path}")
-                    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())