@neoline/hostbridge 2.0.3 → 2.0.4

package/src/server_control_mcp/mock_ssh.py

@@ -1,115 +1,28 @@
 from __future__ import annotations
 
 import asyncio
-import base64
-import binascii
-import json
-import posixpath
-import re
-import shlex
-import sys
-import tempfile
-import threading
-import time
-import uuid
-from collections import deque
 from collections.abc import Awaitable
-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 .client import HostBridgeClient, HostBridgeError
+from .mock_ssh_exec import HostBridgeSSHProcess
+from .mock_ssh_session import HostBridgeSessionFactory
+from .mock_ssh_sftp import HostBridgeSFTPServer
+from .mock_ssh_tunnel import TunnelConnector
+from .secure_files import create_private_file, ensure_private_directory, replace_private_file, validate_private_file
 
 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
-TUNNEL_FRAME_BYTES = 2048
-TUNNEL_READ_BYTES = 32 * 1024
-TUNNEL_MAX_BUFFER_BYTES = 128 * 1024
-
-_REMOTE_TUNNEL_SCRIPT = r"""
-import base64
-import os
-import socket
-import sys
-import threading
-
-host, port, prefix = sys.argv[1], int(sys.argv[2]), sys.argv[3]
-write_lock = threading.Lock()
-
-def emit(kind, payload=b""):
-    suffix = ":" + base64.b64encode(payload).decode("ascii") if payload else ""
-    with write_lock:
-        print(prefix + ":" + kind + suffix, flush=True)
-
-try:
-    connection = socket.create_connection((host, port), timeout=15)
-    connection.settimeout(None)
-except Exception as exc:
-    emit("ERROR", str(exc).encode("utf-8", errors="replace"))
-    raise SystemExit(1)
-
-emit("READY")
-
-def receive():
-    try:
-        while True:
-            data = connection.recv(2048)
-            if not data:
-                emit("EOF")
-                os._exit(0)
-            emit("DATA", data)
-    except Exception as exc:
-        emit("ERROR", str(exc).encode("utf-8", errors="replace"))
-        os._exit(1)
-
-receiver = threading.Thread(target=receive, daemon=True)
-receiver.start()
-try:
-    for raw_line in sys.stdin.buffer:
-        line = raw_line.strip()
-        if line == (prefix + ":EOF").encode("ascii"):
-            connection.shutdown(socket.SHUT_WR)
-            receiver.join()
-            break
-        data_prefix = (prefix + ":DATA:").encode("ascii")
-        if line.startswith(data_prefix):
-            connection.sendall(base64.b64decode(line[len(data_prefix):], validate=True))
-finally:
-    connection.close()
-"""
-
-
-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):
-    def __init__(self, authorized_public_key: str, tcp_forwarder: HostBridgeTCPForwarder | None = None) -> None:
+    def __init__(self, authorized_public_key: str, tunnel_connector: TunnelConnector | None = None) -> None:
         self._authorized_public_key = authorized_public_key.strip()
-        self._tcp_forwarder = tcp_forwarder
+        self._tunnel_connector = tunnel_connector
 
     def begin_auth(self, username: str) -> bool:  # noqa: ARG002
         return True
@@ -130,851 +43,29 @@ class HostBridgeSSHServer(asyncssh.SSHServer):
         orig_host: str,  # noqa: ARG002
         orig_port: int,  # noqa: ARG002
     ) -> Awaitable[asyncssh.SSHTCPSession[bytes]] | bool:
-        if self._tcp_forwarder is None:
+        if self._tunnel_connector is None:
             return False
         try:
-            self._tcp_forwarder.validate_destination(dest_host, dest_port)
+            self._tunnel_connector.validate_destination(dest_host, dest_port)
         except ValueError as exc:
             raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, str(exc)) from exc
-        return self._tcp_forwarder.connect(dest_host, dest_port)
-
-
-class HostBridgeSSHProcess:
-    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[bytes]) -> None:
-        session_id = self._session_factory.open()
-        reusable = False
-        try:
-            reusable = await self._run(process, session_id)
-        except asyncio.CancelledError:
-            self._set_exit_status(process, 255)
-            return
-        finally:
-            self._session_factory.close(session_id, reusable=reusable)
-
-    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 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._client.exec(
-                    session_id,
-                    command,
-                    stdin=stdin_data,
-                    timeout=self._timeout,
-                    output_limit=self._output_limit,
-                    owner="mock-ssh",
-                ),
-            )
-            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)
-            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:
-        waits_for_eof = _command_reads_stdin(command)
-        if process.stdin.at_eof():
-            return b"" if waits_for_eof else None
-        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"" if waits_for_eof else None
-
-        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()
-
-        async def pump_input() -> None:
-            try:
-                while not process.stdin.at_eof():
-                    data = await process.stdin.read(4096)
-                    if not data:
-                        break
-                    await loop.run_in_executor(
-                        None,
-                        partial(
-                            self._client.shell_write,
-                            session_id,
-                            _to_bytes(data),
-                            owner="mock-ssh",
-                        ),
-                    )
-            finally:
-                input_done.set()
-
-        async def pump_output() -> None:
-            try:
-                idle_after_input = 0
-                while not stop.is_set():
-                    result = await loop.run_in_executor(
-                        None,
-                        lambda: self._client.shell_read(
-                            session_id,
-                            timeout=0.2,
-                            max_bytes=4096,
-                            owner="mock-ssh",
-                        ),
-                    )
-                    if result.data:
-                        idle_after_input = 0
-                        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.alive:
-                        stop.set()
-                        break
-            except Exception as exc:
-                print(f"hostbridge mock-ssh shell failed: {exc}", file=sys.stderr, flush=True)
-                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())]
-        await stop.wait()
-        for task in tasks:
-            task.cancel()
-        await asyncio.gather(*tasks, return_exceptions=True)
-        self._set_exit_status(process, 0)
-
-
-@dataclass
-class HostBridgeSFTPHandle:
-    path: str
-    readable: bool = False
-    writable: bool = False
-    data: bytes = b""
-    writes: dict[int, bytes] = field(default_factory=dict)
-    mode: int = 0o644
-
-
-@dataclass(frozen=True, slots=True)
-class _TextCommandResult:
-    output: str
-    exit_code: int | None
-    timed_out: bool
-
-
-class HostBridgeSessionFactory:
-    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:
-        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,
-            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,
-        )
-
-
-class _TunnelFrameReader:
-    def __init__(self, client: HostBridgeClient, session_id: str, prefix: str) -> None:
-        self._client = client
-        self._session_id = session_id
-        self._prefix = prefix.encode("ascii") + b":"
-        self._buffer = bytearray()
-        self._pending: deque[tuple[str, bytes]] = deque()
-
-    async def next(self) -> tuple[str, bytes]:
-        while not self._pending:
-            result = await asyncio.to_thread(
-                self._client.shell_read,
-                self._session_id,
-                timeout=0.2,
-                max_bytes=4096,
-                owner="mock-ssh",
-            )
-            if result.data:
-                self._buffer.extend(result.data)
-                self._parse_lines()
-            if not result.alive and not self._pending:
-                raise HostBridgeError("connection_lost", "remote TCP tunnel closed unexpectedly", retryable=True)
-        return self._pending.popleft()
-
-    def _parse_lines(self) -> None:
-        while b"\n" in self._buffer:
-            raw_line, _, remainder = self._buffer.partition(b"\n")
-            self._buffer = bytearray(remainder)
-            line = raw_line.rstrip(b"\r")
-            marker = line.find(self._prefix)
-            if marker < 0:
-                continue
-            frame = line[marker + len(self._prefix) :]
-            kind, separator, encoded = frame.partition(b":")
-            try:
-                kind_text = kind.decode("ascii")
-            except UnicodeDecodeError:
-                continue
-            if kind_text not in {"READY", "DATA", "EOF", "ERROR", "CLOSED"}:
-                continue
-            payload = b""
-            if separator:
-                try:
-                    payload = base64.b64decode(encoded, validate=True)
-                except (binascii.Error, ValueError) as exc:
-                    raise HostBridgeError("protocol_mismatch", "remote TCP tunnel returned invalid base64") from exc
-            self._pending.append((kind_text, payload))
-        if len(self._buffer) > TUNNEL_MAX_BUFFER_BYTES:
-            marker = self._buffer.rfind(self._prefix)
-            self._buffer = self._buffer[marker:] if marker >= 0 else bytearray()
-
-
-class HostBridgeTCPForwarder:
-    def __init__(
-        self,
-        client: HostBridgeClient,
-        session_factory: HostBridgeSessionFactory,
-        *,
-        connect_timeout: int,
-    ) -> None:
-        self._client = client
-        self._session_factory = session_factory
-        self._connect_timeout = connect_timeout
-
-    @staticmethod
-    def validate_destination(dest_host: str, dest_port: int) -> None:
-        if not isinstance(dest_host, str) or not dest_host.strip() or "\x00" in dest_host or len(dest_host) > 253:
-            raise ValueError("TCP forwarding destination host is invalid")
-        if not isinstance(dest_port, int) or isinstance(dest_port, bool) or not 1 <= dest_port <= 65535:
-            raise ValueError("TCP forwarding destination port must be between 1 and 65535")
-
-    async def connect(self, dest_host: str, dest_port: int) -> asyncssh.SSHTCPSession[bytes]:
-        self.validate_destination(dest_host, dest_port)
-        session_id = await asyncio.to_thread(self._session_factory.open)
-        prefix = f"__HB_TUNNEL_{uuid.uuid4().hex}__"
-        frames = _TunnelFrameReader(self._client, session_id, prefix)
-        try:
-            await self._start_remote(session_id, dest_host, dest_port, prefix)
-            kind, payload = await asyncio.wait_for(frames.next(), timeout=self._connect_timeout)
-            if kind == "ERROR":
-                raise HostBridgeError("connection_lost", payload.decode("utf-8", errors="replace"), retryable=True)
-            if kind != "READY":
-                raise HostBridgeError("protocol_mismatch", f"remote TCP tunnel returned {kind} before READY")
-        except asyncio.CancelledError:
-            await asyncio.to_thread(self._session_factory.close, session_id, reusable=False)
-            raise
-        except Exception as exc:
-            await asyncio.to_thread(self._session_factory.close, session_id, reusable=False)
-            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, str(exc)) from exc
-        return HostBridgeTCPSession(
-            self,
-            session_id=session_id,
-            prefix=prefix,
-            frames=frames,
-            dest_host=dest_host,
-            dest_port=dest_port,
-        )
-
-    async def _start_remote(self, session_id: str, dest_host: str, dest_port: int, prefix: str) -> None:
-        command = (
-            f"__hb_prefix={shlex.quote(prefix)}; stty -echo; "
-            f"python3 -u -c {shlex.quote(_REMOTE_TUNNEL_SCRIPT)} "
-            f"{shlex.quote(dest_host)} {dest_port} \"$__hb_prefix\"; "
-            "__hb_status=$?; stty echo; printf '\n%s:CLOSED\n' \"$__hb_prefix\"\n"
-        )
-        await asyncio.to_thread(
-            self._client.shell_write,
-            session_id,
-            command.encode("utf-8"),
-            owner="mock-ssh",
-        )
-
-    async def _send_data(self, session_id: str, prefix: str, payload: bytes) -> None:
-        lines = [
-            f"{prefix}:DATA:{base64.b64encode(payload[offset : offset + TUNNEL_FRAME_BYTES]).decode('ascii')}\n"
-            for offset in range(0, len(payload), TUNNEL_FRAME_BYTES)
-        ]
-        await self._send_lines(session_id, lines)
-
-    async def _send_lines(self, session_id: str, lines: list[str]) -> None:
-        await asyncio.to_thread(
-            self._client.shell_write,
-            session_id,
-            "".join(lines).encode("ascii"),
-            owner="mock-ssh",
-        )
-
-
-
-_TUNNEL_INPUT_EOF = object()
-
-
-class HostBridgeTCPSession(asyncssh.SSHTCPSession[bytes]):
-    def __init__(
-        self,
-        forwarder: HostBridgeTCPForwarder,
-        *,
-        session_id: str,
-        prefix: str,
-        frames: _TunnelFrameReader,
-        dest_host: str,
-        dest_port: int,
-    ) -> None:
-        self._forwarder = forwarder
-        self._session_id = session_id
-        self._prefix = prefix
-        self._frames = frames
-        self._dest_host = dest_host
-        self._dest_port = dest_port
-        self._channel: asyncssh.SSHTCPChannel[bytes] | None = None
-        self._input: asyncio.Queue[bytes | object] = asyncio.Queue()
-        self._write_ready = asyncio.Event()
-        self._write_ready.set()
-        self._runner: asyncio.Task[None] | None = None
-        self._orphan_cleanup: asyncio.Task[None] | None = None
-        self._released = False
-        self._finished = asyncio.Event()
-
-    def connection_made(self, chan: asyncssh.SSHTCPChannel[bytes]) -> None:
-        self._channel = chan
-
-    def session_started(self) -> None:
-        self._runner = asyncio.create_task(self._run())
-        self._runner.add_done_callback(self._runner_done)
-
-    def data_received(self, data: bytes, datatype: asyncssh.DataType) -> None:  # noqa: ARG002
-        if data:
-            self._input.put_nowait(bytes(data))
-
-    def eof_received(self) -> bool:
-        self._input.put_nowait(_TUNNEL_INPUT_EOF)
-        return True
-
-    def connection_lost(self, exc: Exception | None) -> None:  # noqa: ARG002
-        self.abort()
-
-    def pause_writing(self) -> None:
-        self._write_ready.clear()
-
-    def resume_writing(self) -> None:
-        self._write_ready.set()
-
-    def abort(self) -> None:
-        if self._runner is not None and self._runner is not asyncio.current_task() and not self._runner.done():
-            self._runner.cancel()
-        elif self._runner is None:
-            self._schedule_orphan_cleanup()
-
-    async def wait_finished(self) -> None:
-        await self._finished.wait()
-
-    def _runner_done(self, task: asyncio.Task[None]) -> None:  # noqa: ARG002
-        if not self._finished.is_set():
-            self._schedule_orphan_cleanup()
-
-    def _schedule_orphan_cleanup(self) -> None:
-        if self._orphan_cleanup is None:
-            self._orphan_cleanup = asyncio.create_task(self._finish_orphaned_session())
-
-    async def _finish_orphaned_session(self) -> None:
-        await self._release(reusable=False)
-        if self._channel is not None:
-            with suppress(Exception):
-                self._channel.close()
-        self._finished.set()
-
-    async def _run(self) -> None:
-        input_task = asyncio.create_task(self._pump_input())
-        output_task = asyncio.create_task(self._pump_output())
-        reusable = False
-        try:
-            done, _ = await asyncio.wait({input_task, output_task}, return_when=asyncio.FIRST_COMPLETED)
-            if output_task in done:
-                input_task.cancel()
-                await asyncio.gather(input_task, return_exceptions=True)
-                await output_task
-            else:
-                await input_task
-                await output_task
-            reusable = True
-        except asyncio.CancelledError:
-            raise
-        except Exception as exc:
-            print(
-                f"hostbridge mock-ssh TCP forwarding to {self._dest_host}:{self._dest_port} closed: {exc}",
-                file=sys.stderr,
-                flush=True,
-            )
-        finally:
-            for task in (input_task, output_task):
-                if not task.done():
-                    task.cancel()
-            await asyncio.gather(input_task, output_task, return_exceptions=True)
-            await self._release(reusable=reusable)
-            if self._channel is not None:
-                with suppress(Exception):
-                    self._channel.close()
-            self._finished.set()
-
-    async def _release(self, *, reusable: bool) -> None:
-        if self._released:
-            return
-        self._released = True
-        await asyncio.to_thread(
-            self._forwarder._session_factory.close,
-            self._session_id,
-            reusable=reusable,
-        )
-
-    async def _pump_input(self) -> None:
-        while True:
-            item = await self._input.get()
-            if item is _TUNNEL_INPUT_EOF:
-                await self._forwarder._send_lines(self._session_id, [f"{self._prefix}:EOF\n"])
-                return
-            await self._forwarder._send_data(self._session_id, self._prefix, cast(bytes, item))
-
-    async def _pump_output(self) -> None:
-        if self._channel is None:
-            raise RuntimeError("TCP tunnel channel was not initialized")
-        saw_eof = False
-        while True:
-            kind, payload = await self._frames.next()
-            if kind == "DATA":
-                await self._write_ready.wait()
-                self._channel.write(payload)
-            elif kind == "EOF":
-                self._channel.write_eof()
-                saw_eof = True
-            elif kind == "CLOSED":
-                if not saw_eof:
-                    raise HostBridgeError("protocol_mismatch", "remote TCP tunnel closed without EOF")
-                return
-            elif kind == "ERROR":
-                raise HostBridgeError(
-                    "connection_lost",
-                    payload.decode("utf-8", errors="replace") or "remote TCP tunnel failed",
-                    retryable=True,
-                )
-
-
-class HostBridgeSFTPServer(asyncssh.SFTPServer):
-    def __init__(
-        self,
-        chan: asyncssh.SSHServerChannel,
-        client: HostBridgeClient,
-        session_factory: HostBridgeSessionFactory,
-        timeout: int,
-        max_bytes: int,
-    ) -> None:
-        super().__init__(chan)
-        self._client = client
-        self._session_factory = session_factory
-        self._session_id = session_factory.open()
-        self._timeout = timeout
-        self._max_bytes = max_bytes
-
-    def exit(self) -> None:
-        self._session_factory.close(self._session_id)
-
-    def open(self, path: bytes, pflags: int, attrs: asyncssh.SFTPAttrs) -> HostBridgeSFTPHandle:
-        remote_path = self._path(path)
-        readable = bool(pflags & asyncssh.FXF_READ)
-        writable = bool(pflags & asyncssh.FXF_WRITE)
-        mode = attrs.permissions if attrs.permissions is not None else 0o644
-        data = b""
-        if readable or (writable and not (pflags & asyncssh.FXF_TRUNC) and not (pflags & asyncssh.FXF_CREAT)):
-            try:
-                data = self._download(remote_path)
-            except asyncssh.SFTPNoSuchFile:
-                if not writable:
-                    raise
-                data = b""
-        return HostBridgeSFTPHandle(remote_path, readable=readable, writable=writable, data=data, mode=mode)
-
-    def read(self, file_obj: object, offset: int, size: int) -> bytes:
-        handle = self._handle(file_obj)
-        if not handle.readable:
-            raise asyncssh.SFTPPermissionDenied("file is not open for reading")
-        return handle.data[offset : offset + size]
-
-    def write(self, file_obj: object, offset: int, data: bytes) -> int:
-        handle = self._handle(file_obj)
-        if not handle.writable:
-            raise asyncssh.SFTPPermissionDenied("file is not open for writing")
-        handle.writes[int(offset)] = bytes(data)
-        return len(data)
-
-    def close(self, file_obj: object) -> None:
-        handle = self._handle(file_obj)
-        if handle.writable:
-            self._upload(handle)
-
-    def stat(self, path: bytes) -> asyncssh.SFTPAttrs:
-        return self._stat(self._path(path), follow=True)
-
-    def lstat(self, path: bytes) -> asyncssh.SFTPAttrs:
-        return self._stat(self._path(path), follow=False)
-
-    def fstat(self, file_obj: object) -> asyncssh.SFTPAttrs:
-        handle = self._handle(file_obj)
-        if handle.writable and handle.writes:
-            size = len(handle.data)
-            for offset, chunk in handle.writes.items():
-                size = max(size, offset + len(chunk))
-            return asyncssh.SFTPAttrs(size=size, permissions=handle.mode, type=asyncssh.FILEXFER_TYPE_REGULAR)
-        return self._stat(handle.path, follow=True)
-
-    def fsetstat(self, file_obj: object, attrs: asyncssh.SFTPAttrs) -> None:
-        handle = self._handle(file_obj)
-        if attrs.permissions is not None:
-            handle.mode = attrs.permissions
-        if attrs.atime is not None or attrs.mtime is not None:
-            # Applied after upload in close(); the handle may not exist remotely yet.
-            return
-
-    async def scandir(self, path: bytes):
-        for item in self._listdir(self._path(path)):
-            yield item
-
-    def mkdir(self, path: bytes, attrs: asyncssh.SFTPAttrs) -> None:
-        mode = attrs.permissions if attrs.permissions is not None else 0o755
-        self._checked_run(f"mkdir -p -m {shlex.quote(format(mode & 0o7777, 'o'))} -- {shlex.quote(self._path(path))}")
-
-    def rmdir(self, path: bytes) -> None:
-        self._checked_run(f"rmdir -- {shlex.quote(self._path(path))}")
-
-    def remove(self, path: bytes) -> None:
-        self._checked_run(f"rm -f -- {shlex.quote(self._path(path))}")
-
-    def rename(self, oldpath: bytes, newpath: bytes) -> None:
-        self._checked_run(f"mv -- {shlex.quote(self._path(oldpath))} {shlex.quote(self._path(newpath))}")
-
-    def realpath(self, path: bytes) -> bytes:
-        script = (
-            "import os,sys; "
-            "path=sys.argv[1]; "
-            "parent=os.path.dirname(path) or '/'; "
-            "base=os.path.basename(path); "
-            "print(os.path.join(os.path.realpath(parent), base) if not os.path.exists(path) else os.path.realpath(path))"
-        )
-        command = "python3 -c " + shlex.quote(script) + " " + shlex.quote(self._path(path))
-        result = self._checked_run(command)
-        return (result.output.strip() or self._path(path)).encode("utf-8")
-
-    def setstat(self, path: bytes, attrs: asyncssh.SFTPAttrs) -> None:
-        remote_path = self._path(path)
-        if attrs.permissions is not None:
-            self._checked_run(f"chmod {shlex.quote(format(attrs.permissions & 0o7777, 'o'))} -- {shlex.quote(remote_path)}")
-        if attrs.atime is not None or attrs.mtime is not None:
-            script = "import os,sys; os.utime(sys.argv[1], (int(sys.argv[2]), int(sys.argv[3])))"
-            current = self._stat(remote_path, follow=True)
-            atime = int(attrs.atime if attrs.atime is not None else current.atime or time.time())
-            mtime = int(attrs.mtime if attrs.mtime is not None else current.mtime or time.time())
-            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:
-            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,
-                temporary_path,
-                timeout=self._timeout,
-                owner="mock-ssh",
-            )
-            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
-        finally:
-            if temporary_path is not None:
-                temporary_path.unlink(missing_ok=True)
-
-    def _upload(self, handle: HostBridgeSFTPHandle) -> None:
-        size = len(handle.data)
-        for offset, chunk in handle.writes.items():
-            size = max(size, offset + len(chunk))
-        data = bytearray(handle.data)
-        if len(data) < size:
-            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:
-            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,
-                mode=handle.mode & 0o7777,
-                timeout=self._timeout,
-                owner="mock-ssh",
-            )
-        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"""
-import json, os, stat, sys
-path = sys.argv[1]
-st = os.stat(path) if sys.argv[2] == "1" else os.lstat(path)
-print(json.dumps({
-    "size": st.st_size,
-    "permissions": st.st_mode,
-    "uid": st.st_uid,
-    "gid": st.st_gid,
-    "atime": int(st.st_atime),
-    "mtime": int(st.st_mtime),
-    "type": "dir" if stat.S_ISDIR(st.st_mode) else "link" if stat.S_ISLNK(st.st_mode) else "file",
-}))
-"""
-        result = self._checked_run(f"python3 -c {shlex.quote(script)} {shlex.quote(remote_path)} {'1' if follow else '0'}")
-        return self._attrs(json.loads(result.output))
-
-    def _listdir(self, remote_path: str) -> list[asyncssh.SFTPName]:
-        script = r"""
-import json, os, stat, sys, time
-path = sys.argv[1]
-items = []
-for name in os.listdir(path):
-    full = os.path.join(path, name)
-    st = os.lstat(full)
-    item_type = "dir" if stat.S_ISDIR(st.st_mode) else "link" if stat.S_ISLNK(st.st_mode) else "file"
-    items.append({
-        "filename": name,
-        "longname": name,
-        "attrs": {
-            "size": st.st_size,
-            "permissions": st.st_mode,
-            "uid": st.st_uid,
-            "gid": st.st_gid,
-            "atime": int(st.st_atime),
-            "mtime": int(st.st_mtime),
-            "type": item_type,
-        },
-    })
-print(json.dumps(items))
-"""
-        result = self._checked_run(f"python3 -c {shlex.quote(script)} {shlex.quote(remote_path)}")
-        return [
-            asyncssh.SFTPName(item["filename"], item["longname"], self._attrs(item["attrs"]))
-            for item in json.loads(result.output or "[]")
-        ]
-
-    def _attrs(self, data: dict[str, object]) -> asyncssh.SFTPAttrs:
-        permissions = self._int_value(data.get("permissions"), 0)
-        raw_type = data.get("type")
-        file_type = {
-            "file": asyncssh.FILEXFER_TYPE_REGULAR,
-            "dir": asyncssh.FILEXFER_TYPE_DIRECTORY,
-            "link": asyncssh.FILEXFER_TYPE_SYMLINK,
-        }.get(str(raw_type), asyncssh.FILEXFER_TYPE_UNKNOWN)
-        return asyncssh.SFTPAttrs(
-            type=file_type,
-            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=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._client.exec(
-                self._session_id,
-                command,
-                timeout=self._timeout,
-                owner="mock-ssh",
-            )
-        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_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{error_output}",
-                file=sys.stderr,
-                flush=True,
-            )
-            raise error
-        return _TextCommandResult(output, result.exit_code, result.timed_out)
-
-    def _path(self, path: bytes) -> str:
-        text = path.decode("utf-8", "surrogateescape")
-        if not text:
-            return "."
-        return posixpath.normpath(text)
-
-    def _handle(self, file_obj: object) -> HostBridgeSFTPHandle:
-        if not isinstance(file_obj, HostBridgeSFTPHandle):
-            raise asyncssh.SFTPInvalidHandle("invalid HostBridge SFTP handle")
-        return file_obj
-
-    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:
-            return asyncssh.SFTPNoSuchFile(text)
-        if "permission denied" in lowered:
-            return asyncssh.SFTPPermissionDenied(text)
-        return asyncssh.SFTPFailure(text)
+        return self._tunnel_connector.connect(dest_host, dest_port)
 
 
 def ensure_keypair(key_dir: Path) -> tuple[Path, Path, str]:
-    key_dir.mkdir(parents=True, exist_ok=True)
-    with _suppress_chmod_error(key_dir):
-        key_dir.chmod(0o700)
+    ensure_private_directory(key_dir)
     host_key_path = key_dir / "ssh_host_ed25519_key"
     client_key_path = key_dir / "client_ed25519_key"
-    if not host_key_path.exists():
-        host_key = asyncssh.generate_private_key("ssh-ed25519")
-        host_key_path.write_bytes(host_key.export_private_key())
-    if not client_key_path.exists():
-        client_key = asyncssh.generate_private_key("ssh-ed25519")
-        client_key_path.write_bytes(client_key.export_private_key())
-        (key_dir / "client_ed25519_key.pub").write_bytes(client_key.export_public_key())
-    with _suppress_chmod_error(host_key_path):
-        host_key_path.chmod(0o600)
-    with _suppress_chmod_error(client_key_path):
-        client_key_path.chmod(0o600)
-    public_key = (key_dir / "client_ed25519_key.pub").read_text(encoding="utf-8").strip()
+    for path in (host_key_path, client_key_path):
+        if path.exists():
+            validate_private_file(path)
+            continue
+        private_key = asyncssh.generate_private_key("ssh-ed25519")
+        create_private_file(path, private_key.export_private_key())
+    client_key = asyncssh.import_private_key(client_key_path.read_bytes())
+    public_key_path = key_dir / "client_ed25519_key.pub"
+    replace_private_file(public_key_path, client_key.export_public_key())
+    public_key = public_key_path.read_text(encoding="utf-8").strip()
     return host_key_path, client_key_path, public_key
 
 
@@ -1019,7 +110,9 @@ async def create_server_on_available_port(
             continue
     if port is not None:
         raise RuntimeError(f"port {port} is not available on {listen_host}") from last_error
-    raise RuntimeError(f"no free port found in {DEFAULT_PORT_RANGE.start}-{DEFAULT_PORT_RANGE.stop - 1}") from last_error
+    raise RuntimeError(
+        f"no free port found in {DEFAULT_PORT_RANGE.start}-{DEFAULT_PORT_RANGE.stop - 1}"
+    ) from last_error
 
 
 def install_ssh_config_block(
@@ -1085,13 +178,15 @@ def _remove_managed_block(existing: str, begin: str, end: str) -> str:
     stop += len(end)
     while stop < len(existing) and existing[stop] in "\r\n":
         stop += 1
-    return existing[:start].rstrip() + ("\n" if existing[:start].strip() and existing[stop:].strip() else "") + existing[stop:].lstrip()
+    return (
+        existing[:start].rstrip()
+        + ("\n" if existing[:start].strip() and existing[stop:].strip() else "")
+        + existing[stop:].lstrip()
+    )
 
 
 def _write_text_atomic(path: Path, text: str) -> None:
-    tmp_path = path.with_name(f".{path.name}.hostbridge.tmp")
-    tmp_path.write_text(text, encoding="utf-8")
-    tmp_path.replace(path)
+    replace_private_file(path, text.encode("utf-8"))
 
 
 class _suppress_chmod_error:
@@ -1138,36 +233,40 @@ async def serve_mock_ssh(
     client = HostBridgeClient()
     client.hello()
     session_factory = HostBridgeSessionFactory(client, host_id, connect_timeout)
-    tcp_forwarder = HostBridgeTCPForwarder(client, session_factory, connect_timeout=connect_timeout)
+    tunnel_connector = TunnelConnector(client, session_factory, connect_timeout=connect_timeout)
     server, selected_port = await create_server_on_available_port(
-        server_factory=lambda: HostBridgeSSHServer(public_key, tcp_forwarder),
+        server_factory=lambda: HostBridgeSSHServer(public_key, tunnel_connector),
         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=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, SCP, and TCP forwarding; press Ctrl-C to stop", flush=True)
     try:
+        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, SCP, and TCP forwarding; press Ctrl-C to stop", flush=True)
         await server.wait_closed()
         return 0
     finally:
-        session_factory.shutdown()
+        try:
+            server.close()
+            await server.wait_closed()
+        finally:
+            session_factory.shutdown()
 
 
 def run_mock_ssh(**kwargs: object) -> int: