@neoline/hostbridge 1.4.4 → 2.0.0

package/src/server_control_mcp/client.py

@@ -0,0 +1,468 @@
+"""Typed synchronous client for the HostBridge daemon protocol."""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import hashlib
+import json
+import socket
+import threading
+import uuid
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+
+from .protocol import (
+    MAX_FRAME_BYTES,
+    BinaryFrame,
+    ErrorInfo,
+    FrameFlags,
+    FrameSequenceValidator,
+    ProtocolError,
+    Request,
+    Response,
+    encode_frame,
+    read_frame,
+)
+from .runtime import runtime_paths
+from .transports.base import ShellReadResult, TransferResult
+
+MAX_CONTROL_LINE_BYTES = 4 * 1024 * 1024
+
+
+def _required_text(value: object, field: str) -> str:
+    if not isinstance(value, str) or not value.strip():
+        raise HostBridgeError("protocol_mismatch", f"response field {field} must be a non-empty string")
+    return value
+
+
+def _encode_inline_bytes(value: bytes | None) -> dict[str, str] | None:
+    if value is None:
+        return None
+    if not isinstance(value, bytes):
+        raise TypeError("stdin must be bytes or None")
+    return {"encoding": "base64", "data": base64.b64encode(value).decode("ascii")}
+
+
+def _decode_inline_bytes(value: object, field: str) -> bytes:
+    if not isinstance(value, dict) or value.get("encoding") != "base64" or not isinstance(value.get("data"), str):
+        raise HostBridgeError("protocol_mismatch", f"response field {field} must contain base64 data")
+    try:
+        return base64.b64decode(value["data"], validate=True)
+    except (binascii.Error, ValueError) as exc:
+        raise HostBridgeError("protocol_mismatch", f"response field {field} contains invalid base64") from exc
+
+
+@dataclass(frozen=True, slots=True)
+class SessionInfo:
+    session_id: str
+    host_id: str
+    owner: str | None = None
+
+    @classmethod
+    def from_dict(cls, value: Mapping[str, object]) -> SessionInfo:
+        owner = value.get("owner")
+        if owner is not None and not isinstance(owner, str):
+            raise HostBridgeError("protocol_mismatch", "response field owner must be a string or null")
+        return cls(
+            session_id=_required_text(value.get("session_id"), "session_id"),
+            host_id=_required_text(value.get("host_id"), "host_id"),
+            owner=owner,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class ExecResult:
+    session_id: str
+    exit_code: int | None
+    stdout: bytes
+    stderr: bytes
+    timed_out: bool = False
+
+    @classmethod
+    def from_dict(cls, value: Mapping[str, object]) -> ExecResult:
+        exit_code = value.get("exit_code")
+        if exit_code is not None and (not isinstance(exit_code, int) or isinstance(exit_code, bool)):
+            raise HostBridgeError("protocol_mismatch", "response field exit_code must be an integer or null")
+        timed_out = value.get("timed_out", False)
+        if not isinstance(timed_out, bool):
+            raise HostBridgeError("protocol_mismatch", "response field timed_out must be a boolean")
+        return cls(
+            session_id=_required_text(value.get("session_id"), "session_id"),
+            exit_code=exit_code,
+            stdout=_decode_inline_bytes(value.get("stdout"), "stdout"),
+            stderr=_decode_inline_bytes(value.get("stderr"), "stderr"),
+            timed_out=timed_out,
+        )
+
+
+class HostBridgeError(RuntimeError):
+    def __init__(
+        self,
+        code: str,
+        message: str,
+        *,
+        retryable: bool = False,
+        details: dict[str, object] | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.code = code
+        self.retryable = retryable
+        self.details = details or {}
+
+    @classmethod
+    def from_info(cls, error: ErrorInfo) -> HostBridgeError:
+        return cls(error.code, error.message, retryable=error.retryable, details=error.details)
+
+
+class HostBridgeClient:
+    def __init__(
+        self,
+        socket_path: Path | str | None = None,
+        *,
+        default_timeout: float = 30.0,
+    ) -> None:
+        if default_timeout <= 0:
+            raise ValueError("default_timeout must be positive")
+        self.socket_path = runtime_paths(socket_path).socket
+        self.default_timeout = float(default_timeout)
+
+    def request(
+        self,
+        method: str,
+        params: Mapping[str, object],
+        *,
+        timeout: float | None = None,
+    ) -> dict[str, object]:
+        request = Request.new(method, dict(params))
+        response = self._exchange(request, timeout=self.default_timeout if timeout is None else timeout)
+        if not response.ok:
+            if response.error is None:
+                raise HostBridgeError("protocol_mismatch", "failed response omitted error details")
+            raise HostBridgeError.from_info(response.error)
+        return response.result
+
+    def hello(self) -> dict[str, object]:
+        return self.request("system.hello", {})
+
+    def open_session(self, host_id: str, *, owner: str | None = None) -> SessionInfo:
+        result = self.request("session.open", {"host_id": host_id, "owner": owner})
+        return SessionInfo.from_dict(result)
+
+    def exec(
+        self,
+        session_id: str,
+        command: str,
+        *,
+        stdin: bytes | None = None,
+        timeout: float | None = None,
+        output_limit: int | None = None,
+        owner: str | None = None,
+    ) -> ExecResult:
+        params: dict[str, object] = {
+            "session_id": session_id,
+            "command": command,
+            "stdin": _encode_inline_bytes(stdin),
+        }
+        if timeout is not None:
+            params["timeout"] = timeout
+        if output_limit is not None:
+            params["output_limit"] = output_limit
+        if owner is not None:
+            params["owner"] = owner
+        result = self.request("exec.start", params, timeout=timeout)
+        return ExecResult.from_dict(result)
+
+    def close_session(self, session_id: str, *, owner: str | None = None, force: bool = False) -> None:
+        params: dict[str, object] = {"session_id": session_id, "force": force}
+        if owner is not None:
+            params["owner"] = owner
+        self.request("session.close", params)
+
+    def write_text(
+        self,
+        session_id: str,
+        remote_path: str,
+        content: str,
+        *,
+        mode: int = 0o644,
+        max_bytes: int = 4 * 1024 * 1024,
+        owner: str | None = None,
+    ) -> TransferResult:
+        params: dict[str, object] = {
+            "session_id": session_id,
+            "remote_path": remote_path,
+            "content": content,
+            "mode": mode,
+            "max_bytes": max_bytes,
+        }
+        if owner is not None:
+            params["owner"] = owner
+        result = self.request("file.write_text", params)
+        return self._transfer_result(result)
+
+    def read_text(
+        self,
+        session_id: str,
+        remote_path: str,
+        *,
+        max_bytes: int = 4 * 1024 * 1024,
+        owner: str | None = None,
+    ) -> str:
+        params: dict[str, object] = {"session_id": session_id, "remote_path": remote_path, "max_bytes": max_bytes}
+        if owner is not None:
+            params["owner"] = owner
+        result = self.request("file.read_text", params)
+        content = result.get("content")
+        if not isinstance(content, str):
+            raise HostBridgeError("protocol_mismatch", "text response content must be a string")
+        return content
+
+    def shell_write(self, session_id: str, data: bytes, *, owner: str | None = None) -> int:
+        params: dict[str, object] = {
+            "session_id": session_id,
+            "data": _encode_inline_bytes(data),
+        }
+        if owner is not None:
+            params["owner"] = owner
+        result = self.request("shell.write", params)
+        written = result.get("written")
+        if not isinstance(written, int) or isinstance(written, bool) or written < 0:
+            raise HostBridgeError("protocol_mismatch", "shell write response contains invalid byte count")
+        return written
+
+    def shell_read(
+        self,
+        session_id: str,
+        *,
+        timeout: float = 0.2,
+        max_bytes: int = 4096,
+        owner: str | None = None,
+    ) -> ShellReadResult:
+        params: dict[str, object] = {
+            "session_id": session_id,
+            "timeout": timeout,
+            "max_bytes": max_bytes,
+        }
+        if owner is not None:
+            params["owner"] = owner
+        result = self.request("shell.read", params, timeout=max(self.default_timeout, timeout + 1))
+        alive = result.get("alive")
+        if not isinstance(alive, bool):
+            raise HostBridgeError("protocol_mismatch", "shell read response contains invalid alive flag")
+        return ShellReadResult(_decode_inline_bytes(result.get("data"), "data"), alive)
+
+    def upload_file(
+        self,
+        session_id: str,
+        local_path: Path | str,
+        remote_path: str,
+        *,
+        chunk_size: int = 256 * 1024,
+        mode: int = 0o644,
+        timeout: float | None = None,
+        cancel_event: threading.Event | None = None,
+        owner: str | None = None,
+    ) -> TransferResult:
+        self._validate_chunk_size(chunk_size)
+        source = Path(local_path)
+        expected_size = source.stat().st_size
+        stream_id = uuid.uuid4().hex
+        params: dict[str, object] = {
+            "session_id": session_id,
+            "remote_path": remote_path,
+            "stream_id": stream_id,
+            "mode": mode,
+            "expected_size": expected_size,
+        }
+        if owner is not None:
+            params["owner"] = owner
+        request = Request.new("transfer.upload", params)
+        digest = hashlib.sha256()
+        sequence = 0
+        request_timeout = self.default_timeout if timeout is None else timeout
+        try:
+            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
+                connection.settimeout(request_timeout)
+                connection.connect(str(self.socket_path))
+                connection.sendall(request.to_json().encode("utf-8") + b"\n")
+                with source.open("rb") as local_file:
+                    while True:
+                        if cancel_event is not None and cancel_event.is_set():
+                            payload = json.dumps(
+                                {"code": "cancelled", "message": "upload cancelled by client"},
+                                separators=(",", ":"),
+                            ).encode("utf-8")
+                            connection.sendall(
+                                encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.CANCEL, payload))
+                            )
+                            break
+                        chunk = local_file.read(chunk_size)
+                        if not chunk:
+                            connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.EOF)))
+                            break
+                        digest.update(chunk)
+                        connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.DATA, chunk)))
+                        sequence += 1
+                with connection.makefile("rb") as reader:
+                    response = self._read_response(reader, request)
+        except TimeoutError as exc:
+            raise HostBridgeError("timeout", "HostBridge upload timed out", retryable=True) from exc
+        except OSError as exc:
+            raise HostBridgeError("connection_lost", "HostBridge upload connection failed", retryable=True) from exc
+        result = self._transfer_result(response.result)
+        if result.bytes_transferred != expected_size or result.sha256 != digest.hexdigest():
+            raise HostBridgeError("integrity_error", "HostBridge upload integrity check failed")
+        return result
+
+    def download_file(
+        self,
+        session_id: str,
+        remote_path: str,
+        local_path: Path | str,
+        *,
+        chunk_size: int = 256 * 1024,
+        timeout: float | None = None,
+        owner: str | None = None,
+    ) -> TransferResult:
+        self._validate_chunk_size(chunk_size)
+        destination = Path(local_path)
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        temporary = destination.with_name(f".{destination.name}.hostbridge-{uuid.uuid4().hex}.tmp")
+        stream_id = uuid.uuid4().hex
+        params: dict[str, object] = {
+            "session_id": session_id,
+            "remote_path": remote_path,
+            "stream_id": stream_id,
+            "chunk_size": chunk_size,
+        }
+        if owner is not None:
+            params["owner"] = owner
+        request = Request.new("transfer.download", params)
+        digest = hashlib.sha256()
+        transferred = 0
+        validator = FrameSequenceValidator()
+        request_timeout = self.default_timeout if timeout is None else timeout
+        try:
+            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
+                connection.settimeout(request_timeout)
+                connection.connect(str(self.socket_path))
+                connection.sendall(request.to_json().encode("utf-8") + b"\n")
+                with connection.makefile("rb") as reader:
+                    response = self._read_response(reader, request)
+                    if response.result.get("stream_id") != stream_id:
+                        raise HostBridgeError("protocol_mismatch", "download stream_id does not match")
+                    with temporary.open("wb") as local_file:
+                        while True:
+                            frame = read_frame(reader)
+                            validator.accept(frame)
+                            if frame.stream_id != stream_id:
+                                raise HostBridgeError("protocol_mismatch", "download frame stream_id does not match")
+                            if frame.flags & FrameFlags.CANCEL:
+                                self._raise_cancel_frame(frame)
+                            if frame.flags & FrameFlags.DATA:
+                                local_file.write(frame.payload)
+                                digest.update(frame.payload)
+                                transferred += len(frame.payload)
+                            if frame.flags & FrameFlags.EOF:
+                                metadata = self._frame_metadata(frame.payload)
+                                break
+        except HostBridgeError:
+            temporary.unlink(missing_ok=True)
+            raise
+        except TimeoutError as exc:
+            temporary.unlink(missing_ok=True)
+            raise HostBridgeError("timeout", "HostBridge download timed out", retryable=True) from exc
+        except (OSError, ProtocolError) as exc:
+            temporary.unlink(missing_ok=True)
+            raise HostBridgeError("connection_lost", "HostBridge download connection failed", retryable=True) from exc
+        expected = self._transfer_result(metadata)
+        if expected.bytes_transferred != transferred or expected.sha256 != digest.hexdigest():
+            temporary.unlink(missing_ok=True)
+            raise HostBridgeError("integrity_error", "HostBridge download integrity check failed")
+        temporary.replace(destination)
+        return expected
+
+    @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 <= MAX_FRAME_BYTES:
+            raise ValueError(f"chunk_size must be between 4096 and {MAX_FRAME_BYTES} bytes")
+
+    @staticmethod
+    def _transfer_result(value: Mapping[str, object]) -> TransferResult:
+        byte_count = value.get("bytes_transferred")
+        sha256 = value.get("sha256")
+        if not isinstance(byte_count, int) or isinstance(byte_count, bool) or byte_count < 0:
+            raise HostBridgeError("protocol_mismatch", "transfer response contains invalid byte count")
+        if not isinstance(sha256, str) or len(sha256) != 64:
+            raise HostBridgeError("protocol_mismatch", "transfer response contains invalid SHA-256")
+        return TransferResult(byte_count, sha256.lower())
+
+    def _read_response(self, reader, request: Request) -> Response:
+        raw_response = reader.readline(MAX_CONTROL_LINE_BYTES + 1)
+        if not raw_response:
+            raise HostBridgeError("connection_lost", "HostBridge daemon closed before responding", retryable=True)
+        if len(raw_response) > MAX_CONTROL_LINE_BYTES:
+            raise HostBridgeError("protocol_mismatch", "HostBridge daemon response exceeds the control limit")
+        try:
+            response = Response.from_json(raw_response.decode("utf-8"))
+        except (ProtocolError, UnicodeDecodeError) as exc:
+            raise HostBridgeError("protocol_mismatch", "HostBridge daemon returned an invalid response") from exc
+        if response.request_id != request.request_id:
+            raise HostBridgeError("protocol_mismatch", "HostBridge daemon response request_id does not match")
+        if not response.ok:
+            if response.error is None:
+                raise HostBridgeError("protocol_mismatch", "failed response omitted error details")
+            raise HostBridgeError.from_info(response.error)
+        return response
+
+    @staticmethod
+    def _frame_metadata(payload: bytes) -> dict[str, object]:
+        try:
+            value = json.loads(payload.decode("utf-8"))
+        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+            raise HostBridgeError("protocol_mismatch", "download EOF contains invalid metadata") from exc
+        if not isinstance(value, dict):
+            raise HostBridgeError("protocol_mismatch", "download EOF metadata must be an object")
+        return value
+
+    @staticmethod
+    def _raise_cancel_frame(frame: BinaryFrame) -> None:
+        metadata = HostBridgeClient._frame_metadata(frame.payload)
+        code = metadata.get("code")
+        message = metadata.get("message")
+        raise HostBridgeError(
+            code if isinstance(code, str) else "cancelled",
+            message if isinstance(message, str) else "HostBridge transfer cancelled",
+        )
+
+    def _exchange(self, request: Request, *, timeout: float) -> Response:
+        if timeout <= 0:
+            raise ValueError("timeout must be positive")
+        try:
+            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
+                connection.settimeout(timeout)
+                connection.connect(str(self.socket_path))
+                connection.sendall(request.to_json().encode("utf-8") + b"\n")
+                with connection.makefile("rb") as reader:
+                    raw_response = reader.readline(MAX_CONTROL_LINE_BYTES + 1)
+        except TimeoutError as exc:
+            raise HostBridgeError("timeout", "HostBridge daemon request timed out", retryable=True) from exc
+        except OSError as exc:
+            raise HostBridgeError("connection_lost", "HostBridge daemon connection failed", retryable=True) from exc
+        if not raw_response:
+            raise HostBridgeError("connection_lost", "HostBridge daemon closed before responding", retryable=True)
+        if len(raw_response) > MAX_CONTROL_LINE_BYTES:
+            raise HostBridgeError("protocol_mismatch", "HostBridge daemon response exceeds the control limit")
+        return self._read_response_from_bytes(raw_response, request)
+
+    @staticmethod
+    def _read_response_from_bytes(raw_response: bytes, request: Request) -> Response:
+        try:
+            response = Response.from_json(raw_response.decode("utf-8"))
+        except (ProtocolError, UnicodeDecodeError) as exc:
+            raise HostBridgeError("protocol_mismatch", "HostBridge daemon returned an invalid response") from exc
+        if response.request_id != request.request_id:
+            raise HostBridgeError("protocol_mismatch", "HostBridge daemon response request_id does not match")
+        return response

package/src/server_control_mcp/config.py

@@ -0,0 +1,139 @@
+"""HostBridge schema-v1 configuration and migration."""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+
+SCHEMA_VERSION = 1
+_HOST_ID = re.compile(r"^[A-Za-z0-9_.-]+$")
+_CONNECTION_STYLES = {"script", "ssh", "command"}
+_TRANSPORTS = {"auto", "pty", "ssh"}
+
+
+@dataclass(frozen=True, slots=True)
+class ConfigDocument:
+    path: Path
+    data: dict[str, object]
+
+    @property
+    def hosts(self) -> list[dict[str, object]]:
+        value = self.data["hosts"]
+        assert isinstance(value, list)
+        return value
+
+
+def load_v1_config(path: Path | str) -> ConfigDocument:
+    config_path = Path(path).expanduser()
+    try:
+        payload = json.loads(config_path.read_text(encoding="utf-8"))
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"{config_path} is not valid JSON") from exc
+    validated = validate_v1_config(payload, source=config_path)
+    return ConfigDocument(config_path, validated)
+
+
+def validate_v1_config(payload: object, *, source: Path | str = "configuration") -> dict[str, object]:
+    if not isinstance(payload, dict):
+        raise ValueError(f"{source} must be a JSON object")
+    if "schema_version" not in payload:
+        raise ValueError(f"{source} is missing schema_version")
+    if payload.get("schema_version") != SCHEMA_VERSION:
+        raise ValueError(f"{source} has unsupported schema_version {payload.get('schema_version')!r}")
+    raw_hosts = payload.get("hosts")
+    if not isinstance(raw_hosts, list):
+        raise ValueError(f"{source}.hosts must be a list")
+    hosts: list[dict[str, object]] = []
+    seen: set[str] = set()
+    for index, raw_host in enumerate(raw_hosts):
+        if not isinstance(raw_host, dict):
+            raise ValueError(f"{source}.hosts[{index}] must be an object")
+        host = dict(raw_host)
+        host_id = host.get("id")
+        if not isinstance(host_id, str) or not _HOST_ID.fullmatch(host_id):
+            raise ValueError(f"{source}.hosts[{index}].id is invalid")
+        if host_id in seen:
+            raise ValueError(f"{source} contains duplicate host id {host_id!r}")
+        seen.add(host_id)
+        styles = [name for name in _CONNECTION_STYLES if name in host]
+        if len(styles) != 1:
+            raise ValueError(f"{source}.hosts[{index}] must define exactly one connection style")
+        transport = host.get("transport", "auto")
+        if transport not in _TRANSPORTS:
+            raise ValueError(f"{source}.hosts[{index}].transport must be one of: auto, pty, ssh")
+        host["transport"] = transport
+        _validate_secrets(host.get("secrets", {}), source, index)
+        _validate_limits(host.get("limits"), source, index)
+        hosts.append(host)
+    result = dict(payload)
+    result["schema_version"] = SCHEMA_VERSION
+    result["hosts"] = hosts
+    return result
+
+
+def _validate_secrets(value: object, source: Path | str, index: int) -> None:
+    if value in (None, {}):
+        return
+    if not isinstance(value, dict):
+        raise ValueError(f"{source}.hosts[{index}].secrets must be an object")
+    for name, secret in value.items():
+        if not isinstance(name, str):
+            raise ValueError(f"{source}.hosts[{index}].secrets contains an invalid name")
+        encrypted = isinstance(secret, str) and secret.startswith("gAAAA")
+        environment_ref = isinstance(secret, dict) and set(secret) == {"env"} and isinstance(secret.get("env"), str)
+        if not encrypted and not environment_ref:
+            raise ValueError(f"{source}.hosts[{index}] contains plaintext secret {name!r}")
+
+
+def _validate_limits(value: object, source: Path | str, index: int) -> None:
+    if value is None:
+        return
+    if not isinstance(value, dict):
+        raise ValueError(f"{source}.hosts[{index}].limits must be an object")
+    for name, limit in value.items():
+        if not isinstance(name, str) or not isinstance(limit, int) or isinstance(limit, bool) or limit <= 0:
+            raise ValueError(f"{source}.hosts[{index}].limits values must be positive integers")
+
+
+def migrate_config(source: Path | str, destination: Path | str) -> ConfigDocument:
+    source_path = Path(source).expanduser()
+    destination_path = Path(destination).expanduser()
+    if destination_path.exists():
+        raise FileExistsError(destination_path)
+    try:
+        legacy = json.loads(source_path.read_text(encoding="utf-8"))
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"{source_path} is not valid JSON") from exc
+    if not isinstance(legacy, dict) or not isinstance(legacy.get("hosts"), list):
+        raise ValueError(f"{source_path} must contain a hosts list")
+    migrated_hosts = []
+    for raw_host in legacy["hosts"]:
+        if not isinstance(raw_host, dict):
+            raise ValueError(f"{source_path} contains a non-object host")
+        host = dict(raw_host)
+        host.setdefault("transport", "auto")
+        migrated_hosts.append(host)
+    payload = validate_v1_config(
+        {"schema_version": SCHEMA_VERSION, "hosts": migrated_hosts},
+        source=source_path,
+    )
+    destination_path.parent.mkdir(parents=True, exist_ok=True)
+    encoded = (json.dumps(payload, ensure_ascii=False, indent=2) + "\n").encode("utf-8")
+    fd, temporary_name = tempfile.mkstemp(prefix=f".{destination_path.name}.", dir=destination_path.parent)
+    temporary_path = Path(temporary_name)
+    try:
+        os.fchmod(fd, 0o600)
+        with os.fdopen(fd, "wb") as handle:
+            handle.write(encoded)
+            handle.flush()
+            os.fsync(handle.fileno())
+        temporary_path.replace(destination_path)
+        destination_path.chmod(0o600)
+    except Exception:
+        temporary_path.unlink(missing_ok=True)
+        raise
+    return ConfigDocument(destination_path, payload)