@neoline/hostbridge 2.0.2 → 2.0.4

package/src/server_control_mcp/mock_ssh.py

@@ -1,54 +1,28 @@
 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 collections.abc import Awaitable
 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
-
-
-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) -> None:
+    def __init__(self, authorized_public_key: str, tunnel_connector: TunnelConnector | None = None) -> None:
         self._authorized_public_key = authorized_public_key.strip()
+        self._tunnel_connector = tunnel_connector
 
     def begin_auth(self, username: str) -> bool:  # noqa: ARG002
         return True
@@ -62,562 +36,36 @@ class HostBridgeSSHServer(asyncssh.SSHServer):
     def validate_public_key(self, username: str, key: asyncssh.SSHKey) -> bool:  # noqa: ARG002
         return key.export_public_key().decode("utf-8").strip() == self._authorized_public_key
 
-
-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 HostBridgeSFTPServer(asyncssh.SFTPServer):
-    def __init__(
+    def connection_requested(
         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")
+        dest_host: str,
+        dest_port: int,
+        orig_host: str,  # noqa: ARG002
+        orig_port: int,  # noqa: ARG002
+    ) -> Awaitable[asyncssh.SSHTCPSession[bytes]] | bool:
+        if self._tunnel_connector is None:
+            return False
         try:
-            return int(value)
+            self._tunnel_connector.validate_destination(dest_host, dest_port)
         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)
+            raise asyncssh.ChannelOpenError(asyncssh.OPEN_CONNECT_FAILED, str(exc)) from exc
+        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
 
 
@@ -662,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(
@@ -728,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:
@@ -781,35 +233,40 @@ async def serve_mock_ssh(
     client = HostBridgeClient()
     client.hello()
     session_factory = HostBridgeSessionFactory(client, host_id, 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),
+        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, and SCP; 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: