@neoline/hostbridge 1.4.4 → 2.0.1

package/src/server_control_mcp/mock_ssh.py

@@ -3,19 +3,47 @@ from __future__ import annotations
 import asyncio
 import json
 import posixpath
+import re
 import shlex
 import sys
+import tempfile
+import threading
 import time
 from contextlib import suppress
 from dataclasses import dataclass, field
+from functools import partial
 from pathlib import Path
+from typing import Any, cast
 
 import asyncssh
 
-from .daemon import build_manager
-from .manager import DEFAULT_COMMAND_TIMEOUT, DEFAULT_MAX_TRANSFER_BYTES, RemoteCommandError, SessionManager, _new_token
+from .client import HostBridgeClient, HostBridgeError
 
 DEFAULT_PORT_RANGE = range(2222, 2300)
+DEFAULT_COMMAND_TIMEOUT = 60
+DEFAULT_MAX_TRANSFER_BYTES = 1024 * 1024 * 1024
+EXEC_STDIN_PROBE_TIMEOUT = 0.05
+STDIN_READING_COMMANDS = [
+    re.compile(r"\bcat\s*(?:>|>>)"),
+    re.compile(r"\btee(?:\s|$)"),
+    re.compile(r"(?:^|\s)(?:/bin/)?(?:sh|bash|dash|zsh)\b[^\n;]*\s-s(?:\s|$)"),
+    re.compile(r"(?:^|\s)(?:python[\d.]*|node)(?:\s+\S+)*\s+-($|\s)"),
+    re.compile(r"(?:^|\s)tar\s+(?:-[A-Za-z]*x[A-Za-z]*|[A-Za-z]*x[A-Za-z]*)(?:\s|$)"),
+    re.compile(r"(?:^|\s)tar\s+.*(?:^|\s)f\s*-(?:\s|$)"),
+    re.compile(r"(?:^|\s)dd\s+.*\bof="),
+    re.compile(r"(?:^|\s)scp(?:\s+\S+)*\s+-[A-Za-z]*(?:t|f)[A-Za-z]*(?:\s|$)"),
+]
+MAX_IDLE_BACKEND_SESSIONS = 4
+
+
+def _to_bytes(chunk: str | bytes) -> bytes:
+    if isinstance(chunk, bytes):
+        return chunk
+    return chunk.encode("utf-8")
+
+
+def _command_reads_stdin(command: str | None) -> bool:
+    return bool(command and any(pattern.search(command) for pattern in STDIN_READING_COMMANDS))
 
 
 class HostBridgeSSHServer(asyncssh.SSHServer):
@@ -36,48 +64,95 @@ class HostBridgeSSHServer(asyncssh.SSHServer):
 
 
 class HostBridgeSSHProcess:
-    def __init__(self, manager: SessionManager, session_factory: "HostBridgeSessionFactory", timeout: int, output_limit: int) -> None:
-        self._manager = manager
+    def __init__(self, client: HostBridgeClient, session_factory: HostBridgeSessionFactory, timeout: int, output_limit: int) -> None:
+        self._client = client
         self._session_factory = session_factory
         self._timeout = timeout
         self._output_limit = output_limit
 
-    async def __call__(self, process: asyncssh.SSHServerProcess[str]) -> None:
+    async def __call__(self, process: asyncssh.SSHServerProcess[bytes]) -> None:
         session_id = self._session_factory.open()
+        reusable = False
         try:
-            await self._run(process, session_id)
+            reusable = await self._run(process, session_id)
+        except asyncio.CancelledError:
+            self._set_exit_status(process, 255)
+            return
         finally:
-            self._session_factory.close(session_id)
+            self._session_factory.close(session_id, reusable=reusable)
 
-    async def _run(self, process: asyncssh.SSHServerProcess[str], session_id: str) -> None:
+    async def _run(self, process: asyncssh.SSHServerProcess[bytes], session_id: str) -> bool:
         command = process.command
         if not command:
             await self._run_shell(process, session_id)
-            return
+            return False
 
         loop = asyncio.get_running_loop()
         try:
+            stdin_data = await self._read_exec_stdin(process, command)
             result = await loop.run_in_executor(
                 None,
-                lambda: self._manager.run_command(
+                lambda: self._client.exec(
                     session_id,
                     command,
+                    stdin=stdin_data,
                     timeout=self._timeout,
-                    token_factory=_new_token,
                     output_limit=self._output_limit,
+                    owner="mock-ssh",
                 ),
             )
-            if result.output:
-                process.stdout.write(result.output)
-                if not result.output.endswith("\n"):
-                    process.stdout.write("\n")
-            process.exit(124 if result.timed_out else int(result.exit_code or 0))
+            if result.stdout:
+                process.stdout.write(result.stdout)
+            if result.stderr:
+                process.stderr.write(result.stderr)
+            self._set_exit_status(process, 124 if result.timed_out else int(result.exit_code or 0))
+            return not result.timed_out
         except Exception as exc:
             print(f"hostbridge mock-ssh exec failed: {exc}", file=sys.stderr, flush=True)
-            process.stderr.write(f"hostbridge mock-ssh error: {exc}\n")
-            process.exit(1)
-
-    async def _run_shell(self, process: asyncssh.SSHServerProcess[str], session_id: str) -> None:
+            with suppress(Exception):
+                process.stderr.write(f"hostbridge mock-ssh error: {exc}\n".encode("utf-8", errors="replace"))
+            status = 124 if isinstance(exc, HostBridgeError) and exc.code == "timeout" else 1
+            self._set_exit_status(process, status)
+            return False
+
+    @staticmethod
+    def _set_exit_status(process: asyncssh.SSHServerProcess[bytes], status: int) -> None:
+        with suppress(Exception):
+            process.exit(status)
+
+    async def _read_exec_stdin(self, process: asyncssh.SSHServerProcess[bytes], command: str) -> bytes | None:
+        if process.stdin.at_eof():
+            return b""
+        waits_for_eof = _command_reads_stdin(command)
+        try:
+            first_chunk = await asyncio.wait_for(
+                process.stdin.read(4096),
+                timeout=self._timeout if waits_for_eof else EXEC_STDIN_PROBE_TIMEOUT,
+            )
+        except TimeoutError:
+            if waits_for_eof:
+                raise HostBridgeError("timeout", "timed out waiting for mock-ssh exec stdin") from None
+            return None
+        if not first_chunk:
+            return b""
+
+        first_chunk_bytes = _to_bytes(first_chunk)
+        total = len(first_chunk_bytes)
+        if total > self._output_limit:
+            raise HostBridgeError("resource_limit", f"mock-ssh exec stdin exceeds limit ({self._output_limit} bytes)")
+        chunks = [first_chunk_bytes]
+        while not process.stdin.at_eof():
+            chunk = await asyncio.wait_for(process.stdin.read(4096), timeout=self._timeout)
+            if not chunk:
+                break
+            chunk_bytes = _to_bytes(chunk)
+            total += len(chunk_bytes)
+            if total > self._output_limit:
+                raise HostBridgeError("resource_limit", f"mock-ssh exec stdin exceeds limit ({self._output_limit} bytes)")
+            chunks.append(chunk_bytes)
+        return b"".join(chunks)
+
+    async def _run_shell(self, process: asyncssh.SSHServerProcess[bytes], session_id: str) -> None:
         loop = asyncio.get_running_loop()
         stop = asyncio.Event()
         input_done = asyncio.Event()
@@ -90,7 +165,12 @@ class HostBridgeSSHProcess:
                         break
                     await loop.run_in_executor(
                         None,
-                        lambda chunk=data: self._manager.send(session_id, chunk, owner_agent_id="mock-ssh"),
+                        partial(
+                            self._client.shell_write,
+                            session_id,
+                            _to_bytes(data),
+                            owner="mock-ssh",
+                        ),
                     )
             finally:
                 input_done.set()
@@ -101,28 +181,28 @@ class HostBridgeSSHProcess:
                 while not stop.is_set():
                     result = await loop.run_in_executor(
                         None,
-                        lambda: self._manager.read(
+                        lambda: self._client.shell_read(
                             session_id,
                             timeout=0.2,
-                            max_chars=4096,
-                            owner_agent_id="mock-ssh",
+                            max_bytes=4096,
+                            owner="mock-ssh",
                         ),
                     )
-                    output = str(result.get("output") or "")
-                    if output:
+                    if result.data:
                         idle_after_input = 0
-                        process.stdout.write(output)
+                        process.stdout.write(result.data)
                     elif input_done.is_set():
                         idle_after_input += 1
                         if idle_after_input >= 5:
                             stop.set()
                             break
-                    if not result.get("alive", True):
+                    if not result.alive:
                         stop.set()
                         break
             except Exception as exc:
                 print(f"hostbridge mock-ssh shell failed: {exc}", file=sys.stderr, flush=True)
-                process.stderr.write(f"hostbridge mock-ssh shell error: {exc}\n")
+                with suppress(Exception):
+                    process.stderr.write(f"hostbridge mock-ssh shell error: {exc}\n".encode("utf-8", errors="replace"))
                 stop.set()
 
         tasks = [asyncio.create_task(pump_input()), asyncio.create_task(pump_output())]
@@ -130,7 +210,7 @@ class HostBridgeSSHProcess:
         for task in tasks:
             task.cancel()
         await asyncio.gather(*tasks, return_exceptions=True)
-        process.exit(0)
+        self._set_exit_status(process, 0)
 
 
 @dataclass
@@ -143,36 +223,99 @@ class HostBridgeSFTPHandle:
     mode: int = 0o644
 
 
+@dataclass(frozen=True, slots=True)
+class _TextCommandResult:
+    output: str
+    exit_code: int | None
+    timed_out: bool
+
+
 class HostBridgeSessionFactory:
-    def __init__(self, manager: SessionManager, host_id: str, connect_timeout: int) -> None:
-        self._manager = manager
+    def __init__(self, client: HostBridgeClient, host_id: str, connect_timeout: int = 60) -> None:
+        self._client = client
         self._host_id = host_id
         self._connect_timeout = connect_timeout
+        self._lock = threading.Lock()
+        self._idle: list[str] = []
+        self._leased: set[str] = set()
+        self._shutdown = False
 
     def open(self) -> str:
-        session = self._manager.open_session(
+        with self._lock:
+            if self._shutdown:
+                raise RuntimeError("mock-ssh session factory is shut down")
+            if self._idle:
+                session_id = self._idle.pop()
+                self._leased.add(session_id)
+                return session_id
+        session = self._client.open_session(
             self._host_id,
-            connect_timeout=self._connect_timeout,
-            owner_agent_id="mock-ssh",
+            owner="mock-ssh",
+            timeout=self._connect_timeout + 5,
+        )
+        session_id = session.session_id
+        with self._lock:
+            if not self._shutdown:
+                self._leased.add(session_id)
+                return session_id
+        self._close_backend(session_id)
+        raise RuntimeError("mock-ssh session factory is shut down")
+
+    def close(self, session_id: str, *, reusable: bool = True) -> None:
+        should_close = False
+        with self._lock:
+            if session_id not in self._leased:
+                return
+            self._leased.remove(session_id)
+            if reusable and not self._shutdown and len(self._idle) < MAX_IDLE_BACKEND_SESSIONS:
+                self._idle.append(session_id)
+            else:
+                should_close = True
+        if not should_close:
+            return
+        self._close_backend(session_id)
+
+    def shutdown(self) -> None:
+        with self._lock:
+            self._shutdown = True
+            session_ids = [*self._idle, *self._leased]
+            self._idle.clear()
+            self._leased.clear()
+        for session_id in session_ids:
+            self._close_backend(session_id)
+
+    def _close_backend(self, session_id: str) -> None:
+        last_error: Exception | None = None
+        for attempt in range(2):
+            try:
+                self._client.close_session(session_id, owner="mock-ssh", force=True)
+                return
+            except HostBridgeError as exc:
+                if exc.code == "not_found":
+                    return
+                last_error = exc
+            except Exception as exc:
+                last_error = exc
+            if attempt == 0:
+                time.sleep(0.05)
+        print(
+            f"hostbridge mock-ssh failed to close backend session {session_id}: {last_error}",
+            file=sys.stderr,
+            flush=True,
         )
-        return session.id
-
-    def close(self, session_id: str) -> None:
-        with _suppress_error():
-            self._manager.close_session(session_id, owner_agent_id="mock-ssh", force=True)
 
 
 class HostBridgeSFTPServer(asyncssh.SFTPServer):
     def __init__(
         self,
         chan: asyncssh.SSHServerChannel,
-        manager: SessionManager,
+        client: HostBridgeClient,
         session_factory: HostBridgeSessionFactory,
         timeout: int,
         max_bytes: int,
     ) -> None:
         super().__init__(chan)
-        self._manager = manager
+        self._client = client
         self._session_factory = session_factory
         self._session_id = session_factory.open()
         self._timeout = timeout
@@ -278,17 +421,25 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
             self._checked_run(f"python3 -c {shlex.quote(script)} {shlex.quote(remote_path)} {atime} {mtime}")
 
     def _download(self, remote_path: str) -> bytes:
+        temporary_path: Path | None = None
         try:
-            result = self._manager.file_read_bytes(
+            with tempfile.NamedTemporaryFile(prefix="hostbridge-sftp-", delete=False) as temporary:
+                temporary_path = Path(temporary.name)
+            result = self._client.download_file(
                 self._session_id,
                 remote_path,
-                max_bytes=self._max_bytes,
-                token_factory=_new_token,
-                owner_agent_id="mock-ssh",
+                temporary_path,
+                timeout=self._timeout,
+                owner="mock-ssh",
             )
-        except RemoteCommandError as exc:
+            if result.bytes_transferred > self._max_bytes:
+                raise HostBridgeError("resource_limit", f"file exceeds mock-ssh limit ({self._max_bytes} bytes)")
+            return temporary_path.read_bytes()
+        except HostBridgeError as exc:
             raise self._sftp_error(exc) from exc
-        return bytes(result["data_bytes"])
+        finally:
+            if temporary_path is not None:
+                temporary_path.unlink(missing_ok=True)
 
     def _upload(self, handle: HostBridgeSFTPHandle) -> None:
         size = len(handle.data)
@@ -299,19 +450,27 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
             data.extend(b"\x00" * (size - len(data)))
         for offset, chunk in sorted(handle.writes.items()):
             data[offset : offset + len(chunk)] = chunk
+        if len(data) > self._max_bytes:
+            raise asyncssh.SFTPFailure(f"file exceeds mock-ssh limit ({self._max_bytes} bytes)")
+        temporary_path: Path | None = None
         try:
-            self._manager.file_write_bytes(
+            with tempfile.NamedTemporaryFile(prefix="hostbridge-sftp-", delete=False) as temporary:
+                temporary.write(data)
+                temporary_path = Path(temporary.name)
+            self._client.upload_file(
                 self._session_id,
+                temporary_path,
                 handle.path,
-                bytes(data),
-                mode=format(handle.mode & 0o7777, "04o"),
-                max_bytes=self._max_bytes,
-                token_factory=_new_token,
-                owner_agent_id="mock-ssh",
+                mode=handle.mode & 0o7777,
+                timeout=self._timeout,
+                owner="mock-ssh",
             )
-        except RemoteCommandError as exc:
+        except HostBridgeError as exc:
             print(f"hostbridge mock-ssh sftp upload failed: {exc}", file=sys.stderr, flush=True)
             raise self._sftp_error(exc) from exc
+        finally:
+            if temporary_path is not None:
+                temporary_path.unlink(missing_ok=True)
 
     def _stat(self, remote_path: str, *, follow: bool) -> asyncssh.SFTPAttrs:
         script = r"""
@@ -362,7 +521,7 @@ print(json.dumps(items))
         ]
 
     def _attrs(self, data: dict[str, object]) -> asyncssh.SFTPAttrs:
-        permissions = int(data.get("permissions") or 0)
+        permissions = self._int_value(data.get("permissions"), 0)
         raw_type = data.get("type")
         file_type = {
             "file": asyncssh.FILEXFER_TYPE_REGULAR,
@@ -371,39 +530,54 @@ print(json.dumps(items))
         }.get(str(raw_type), asyncssh.FILEXFER_TYPE_UNKNOWN)
         return asyncssh.SFTPAttrs(
             type=file_type,
-            size=int(data.get("size") or 0),
-            uid=int(data.get("uid") or 0),
-            gid=int(data.get("gid") or 0),
+            size=self._int_value(data.get("size"), 0),
+            uid=self._int_value(data.get("uid"), 0),
+            gid=self._int_value(data.get("gid"), 0),
             permissions=permissions,
-            atime=int(data.get("atime") or time.time()),
-            mtime=int(data.get("mtime") or time.time()),
+            atime=self._int_value(data.get("atime"), int(time.time())),
+            mtime=self._int_value(data.get("mtime"), int(time.time())),
         )
 
+    @staticmethod
+    def _int_value(value: object, default: int) -> int:
+        if value is None:
+            return default
+        if isinstance(value, bool) or not isinstance(value, int | float | str):
+            raise asyncssh.SFTPFailure("remote metadata contains a non-numeric field")
+        try:
+            return int(value)
+        except ValueError as exc:
+            raise asyncssh.SFTPFailure("remote metadata contains an invalid numeric field") from exc
+
     def _checked_run(self, command: str):
         try:
-            result = self._manager.run_command(
+            result = self._client.exec(
                 self._session_id,
                 command,
                 timeout=self._timeout,
-                token_factory=_new_token,
-                owner_agent_id="mock-ssh",
+                owner="mock-ssh",
             )
-        except RemoteCommandError as exc:
+        except HostBridgeError as exc:
             raise self._sftp_error(exc) from exc
+        output = result.stdout.decode("utf-8", errors="replace")
+        stderr = result.stderr.decode("utf-8", errors="replace")
         if result.timed_out:
             print(f"hostbridge mock-ssh sftp command timed out: {command}", file=sys.stderr, flush=True)
             raise asyncssh.SFTPFailure("remote command timed out")
         if result.exit_code not in (0, None):
-            error = self._sftp_error(RemoteCommandError(result.output or f"remote command exited {result.exit_code}"))
+            error_output = "\n".join(part for part in (output, stderr) if part).strip()
+            error = self._sftp_error(
+                HostBridgeError("remote_error", error_output or f"remote command exited {result.exit_code}")
+            )
             if isinstance(error, asyncssh.SFTPNoSuchFile):
                 raise error
             print(
-                f"hostbridge mock-ssh sftp command failed exit={result.exit_code}: {command}\n{result.output}",
+                f"hostbridge mock-ssh sftp command failed exit={result.exit_code}: {command}\n{error_output}",
                 file=sys.stderr,
                 flush=True,
             )
             raise error
-        return result
+        return _TextCommandResult(output, result.exit_code, result.timed_out)
 
     def _path(self, path: bytes) -> str:
         text = path.decode("utf-8", "surrogateescape")
@@ -416,7 +590,7 @@ print(json.dumps(items))
             raise asyncssh.SFTPInvalidHandle("invalid HostBridge SFTP handle")
         return file_obj
 
-    def _sftp_error(self, exc: RemoteCommandError) -> Exception:
+    def _sftp_error(self, exc: HostBridgeError) -> Exception:
         text = str(exc)
         lowered = text.lower()
         if "no such file" in lowered or "not found" in lowered or "filenotfounderror" in lowered:
@@ -480,6 +654,7 @@ async def create_server_on_available_port(
                 process_factory=process_factory,
                 sftp_factory=sftp_factory,
                 allow_scp=True,
+                encoding=None,
             )
             return server, int(candidate)
         except OSError as exc:
@@ -595,49 +770,54 @@ async def serve_mock_ssh(
     command_timeout: int = DEFAULT_COMMAND_TIMEOUT,
     output_limit: int = DEFAULT_MAX_TRANSFER_BYTES,
 ) -> int:
+    if config_path is not None:
+        raise HostBridgeError(
+            "invalid_request",
+            "mock-ssh no longer accepts a separate config path; start the v1 daemon with the intended config",
+        )
     key_dir = key_dir or default_key_dir(host_id)
     ssh_host_alias = ssh_host_alias or default_ssh_host_alias(host_id)
     host_key_path, client_key_path, public_key = ensure_keypair(key_dir)
-    manager = build_manager(config_path)
-    session_factory = HostBridgeSessionFactory(manager, host_id, connect_timeout)
-    try:
-        server, selected_port = await create_server_on_available_port(
-            server_factory=lambda: HostBridgeSSHServer(public_key),
+    client = HostBridgeClient()
+    client.hello()
+    session_factory = HostBridgeSessionFactory(client, host_id, connect_timeout)
+    server, selected_port = await create_server_on_available_port(
+        server_factory=lambda: HostBridgeSSHServer(public_key),
+        listen_host=listen_host,
+        port=port,
+        server_host_keys=[str(host_key_path)],
+        process_factory=HostBridgeSSHProcess(client, session_factory, command_timeout, output_limit),
+        sftp_factory=lambda chan: HostBridgeSFTPServer(chan, client, session_factory, command_timeout, output_limit),
+    )
+    installed_config_path: Path | None = None
+    if install_ssh_config:
+        installed_config_path = install_ssh_config_block(
+            host_alias=ssh_host_alias,
             listen_host=listen_host,
-            port=port,
-            server_host_keys=[str(host_key_path)],
-            process_factory=HostBridgeSSHProcess(manager, session_factory, command_timeout, output_limit),
-            sftp_factory=lambda chan: HostBridgeSFTPServer(chan, manager, session_factory, command_timeout, output_limit),
+            port=selected_port,
+            client_key_path=client_key_path,
+            config_path=ssh_config_path,
         )
-        installed_config_path: Path | None = None
-        if install_ssh_config:
-            installed_config_path = install_ssh_config_block(
-                host_alias=ssh_host_alias,
-                listen_host=listen_host,
-                port=selected_port,
-                client_key_path=client_key_path,
-                config_path=ssh_config_path,
-            )
-        print(f"hostbridge mock-ssh forwarding host '{host_id}'", flush=True)
-        print(f"listening on ssh://{listen_host}:{selected_port}", flush=True)
-        print(f"ssh host alias: {ssh_host_alias}", flush=True)
-        if installed_config_path is not None:
-            print(f"ssh config: {installed_config_path}", flush=True)
-        print(f"identity file: {client_key_path}", flush=True)
-        print("supports SSH exec, shell, SFTP, and SCP; press Ctrl-C to stop", flush=True)
+    print(f"hostbridge mock-ssh forwarding host '{host_id}'", flush=True)
+    print(f"listening on ssh://{listen_host}:{selected_port}", flush=True)
+    print(f"ssh host alias: {ssh_host_alias}", flush=True)
+    if installed_config_path is not None:
+        print(f"ssh config: {installed_config_path}", flush=True)
+    print(f"identity file: {client_key_path}", flush=True)
+    print("supports SSH exec, shell, SFTP, and SCP; press Ctrl-C to stop", flush=True)
+    try:
         await server.wait_closed()
+        return 0
     finally:
-        manager.close_all_sessions()
-        manager.stop_idle_reaper()
-    return 0
+        session_factory.shutdown()
 
 
 def run_mock_ssh(**kwargs: object) -> int:
     try:
-        return asyncio.run(serve_mock_ssh(**kwargs))
+        return asyncio.run(cast(Any, serve_mock_ssh)(**kwargs))
     except KeyboardInterrupt:
         return 130
-    except RemoteCommandError as exc:
+    except HostBridgeError as exc:
         print(f"hostbridge mock-ssh failed: {exc}")
         return 1
 

package/src/server_control_mcp/policy.py

@@ -167,6 +167,7 @@ def audit_event(
     event: str,
     session_id: str | None = None,
     host_id: str | None = None,
+    owner_agent_id: str | None = None,
     command: str | None = None,
     task_id: str | None = None,
     outcome: str | None = None,
@@ -177,6 +178,7 @@ def audit_event(
     for key, value in {
         "session_id": session_id,
         "host_id": host_id,
+        "owner_agent_id": owner_agent_id,
         "command": command,
         "task_id": task_id,
         "outcome": outcome,