@neoline/hostbridge 2.0.3 → 2.0.4

package/src/server_control_mcp/mock_ssh_sftp.py

@@ -0,0 +1,506 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import posixpath
+import shlex
+import sys
+import tempfile
+import threading
+import time
+from contextlib import suppress
+from dataclasses import dataclass
+from pathlib import Path
+
+import asyncssh
+
+from .client import HostBridgeClient, HostBridgeError
+from .mock_ssh_session import HostBridgeSessionFactory
+
+
+@dataclass
+class HostBridgeSFTPHandle:
+    path: str
+    temporary_path: Path
+    readable: bool = False
+    writable: bool = False
+    append: bool = False
+    mode: int = 0o644
+    size: int = 0
+    dirty: bool = False
+    atime: int | None = None
+    mtime: int | None = None
+    closed: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class _TextCommandResult:
+    output: str
+    exit_code: int | None
+    timed_out: bool
+
+
+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: str | None = None
+        self._session_lock = threading.Lock()
+        self._operation_lock = asyncio.Lock()
+        self._timeout = timeout
+        self._max_bytes = max_bytes
+        self._handles: list[HostBridgeSFTPHandle] = []
+
+    async def _run_blocking(self, function, /, *args, **kwargs):
+        async with self._operation_lock:
+            worker = asyncio.create_task(asyncio.to_thread(function, *args, **kwargs))
+            try:
+                return await asyncio.shield(worker)
+            except asyncio.CancelledError:
+                while not worker.done():
+                    try:
+                        await asyncio.shield(worker)
+                    except asyncio.CancelledError:
+                        continue
+                    except BaseException:
+                        break
+                with suppress(BaseException):
+                    worker.result()
+                raise
+
+    async def exit(self) -> None:
+        await self._run_blocking(self._exit_sync)
+
+    def _exit_sync(self) -> None:
+        for handle in list(self._handles):
+            with suppress(Exception):
+                self._close_sync(handle)
+        if self._session_id is not None:
+            self._session_factory.close(self._session_id, reusable=False)
+            self._session_id = None
+
+    async def open(self, path: bytes, pflags: int, attrs: asyncssh.SFTPAttrs) -> HostBridgeSFTPHandle:
+        return await self._run_blocking(self._open_sync, path, pflags, attrs)
+
+    def _open_sync(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)
+        create = bool(pflags & asyncssh.FXF_CREAT)
+        truncate = bool(pflags & asyncssh.FXF_TRUNC)
+        exclusive = bool(pflags & asyncssh.FXF_EXCL)
+        append = bool(pflags & asyncssh.FXF_APPEND)
+        if not readable and not writable:
+            raise asyncssh.SFTPFailure("file must be opened for reading or writing")
+        if (truncate or append) and not writable:
+            raise asyncssh.SFTPPermissionDenied("write flags require write access")
+        if create:
+            create_mode = 0o644 if attrs.permissions is None else attrs.permissions
+            self._create_remote(remote_path, mode=create_mode, exclusive=exclusive)
+        remote_attrs = self._stat(remote_path, follow=True)
+        if remote_attrs is not None and remote_attrs.type == asyncssh.FILEXFER_TYPE_DIRECTORY:
+            raise asyncssh.SFTPFileIsADirectory(remote_path)
+        mode = attrs.permissions if attrs.permissions is not None else 0o644
+        if remote_attrs is not None and remote_attrs.permissions is not None and attrs.permissions is None:
+            mode = remote_attrs.permissions
+        temporary_path = self._new_temporary_path()
+        try:
+            if truncate:
+                self._truncate_remote(remote_path)
+                temporary_path.touch(mode=0o600)
+            else:
+                self._download_to(remote_path, temporary_path)
+            size = temporary_path.stat().st_size
+            handle = HostBridgeSFTPHandle(
+                remote_path,
+                temporary_path,
+                readable=readable,
+                writable=writable,
+                append=append,
+                mode=mode,
+                size=size,
+                dirty=False,
+            )
+            self._handles.append(handle)
+            return handle
+        except BaseException:
+            temporary_path.unlink(missing_ok=True)
+            raise
+
+    async def read(self, file_obj: object, offset: int, size: int) -> bytes:
+        return await self._run_blocking(self._read_sync, file_obj, offset, size)
+
+    def _read_sync(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")
+        if offset < 0 or size < 0:
+            raise asyncssh.SFTPFailure("read offset and size must be non-negative")
+        with handle.temporary_path.open("rb") as stream:
+            stream.seek(offset)
+            return stream.read(size)
+
+    async def write(self, file_obj: object, offset: int, data: bytes) -> int:
+        return await self._run_blocking(self._write_sync, file_obj, offset, data)
+
+    def _write_sync(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")
+        if offset < 0:
+            raise asyncssh.SFTPFailure("write offset must be non-negative")
+        payload = bytes(data)
+        write_offset = handle.size if handle.append else int(offset)
+        final_size = max(handle.size, write_offset + len(payload))
+        if final_size > self._max_bytes:
+            raise asyncssh.SFTPFailure(f"file exceeds mock-ssh limit ({self._max_bytes} bytes)")
+        with handle.temporary_path.open("r+b") as stream:
+            stream.seek(write_offset)
+            stream.write(payload)
+        handle.size = final_size
+        handle.dirty = True
+        return len(payload)
+
+    async def close(self, file_obj: object) -> None:
+        await self._run_blocking(self._close_sync, file_obj)
+
+    def _close_sync(self, file_obj: object) -> None:
+        handle = self._handle(file_obj)
+        if handle.closed:
+            return
+        try:
+            if handle.writable and handle.dirty:
+                self._upload(handle)
+            elif handle.atime is not None or handle.mtime is not None:
+                self._apply_times(handle)
+        finally:
+            handle.closed = True
+            handle.temporary_path.unlink(missing_ok=True)
+            with suppress(ValueError):
+                self._handles.remove(handle)
+
+    async def stat(self, path: bytes) -> asyncssh.SFTPAttrs:
+        return await self._run_blocking(self._stat, self._path(path), follow=True)
+
+    async def lstat(self, path: bytes) -> asyncssh.SFTPAttrs:
+        return await self._run_blocking(self._stat, self._path(path), follow=False)
+
+    async def fstat(self, file_obj: object) -> asyncssh.SFTPAttrs:
+        return await self._run_blocking(self._fstat_sync, file_obj)
+
+    def _fstat_sync(self, file_obj: object) -> asyncssh.SFTPAttrs:
+        handle = self._handle(file_obj)
+        return asyncssh.SFTPAttrs(
+            size=handle.size,
+            permissions=handle.mode,
+            type=asyncssh.FILEXFER_TYPE_REGULAR,
+        )
+
+    async def fsetstat(self, file_obj: object, attrs: asyncssh.SFTPAttrs) -> None:
+        await self._run_blocking(self._fsetstat_sync, file_obj, attrs)
+
+    def _fsetstat_sync(self, file_obj: object, attrs: asyncssh.SFTPAttrs) -> None:
+        handle = self._handle(file_obj)
+        if not handle.writable:
+            raise asyncssh.SFTPPermissionDenied("file is not open for writing")
+        if attrs.permissions is not None:
+            handle.mode = attrs.permissions
+            handle.dirty = True
+        if attrs.size is not None:
+            if attrs.size < 0 or attrs.size > self._max_bytes:
+                raise asyncssh.SFTPFailure(f"file exceeds mock-ssh limit ({self._max_bytes} bytes)")
+            with handle.temporary_path.open("r+b") as stream:
+                stream.truncate(attrs.size)
+            handle.size = attrs.size
+            handle.dirty = True
+        handle.atime = attrs.atime if attrs.atime is not None else handle.atime
+        handle.mtime = attrs.mtime if attrs.mtime is not None else handle.mtime
+
+    async def scandir(self, path: bytes):
+        items = await self._run_blocking(self._listdir, self._path(path))
+        for item in items:
+            yield item
+
+    async def mkdir(self, path: bytes, attrs: asyncssh.SFTPAttrs) -> None:
+        await self._run_blocking(self._mkdir_sync, path, attrs)
+
+    def _mkdir_sync(self, path: bytes, attrs: asyncssh.SFTPAttrs) -> None:
+        mode = attrs.permissions if attrs.permissions is not None else 0o755
+        self._checked_run(f"mkdir -m {shlex.quote(format(mode & 0o7777, 'o'))} -- {shlex.quote(self._path(path))}")
+
+    async def rmdir(self, path: bytes) -> None:
+        await self._run_blocking(self._rmdir_sync, path)
+
+    def _rmdir_sync(self, path: bytes) -> None:
+        self._checked_run(f"rmdir -- {shlex.quote(self._path(path))}")
+
+    async def remove(self, path: bytes) -> None:
+        await self._run_blocking(self._remove_sync, path)
+
+    def _remove_sync(self, path: bytes) -> None:
+        self._checked_run(f"rm -- {shlex.quote(self._path(path))}")
+
+    async def rename(self, oldpath: bytes, newpath: bytes) -> None:
+        await self._run_blocking(self._rename_sync, oldpath, newpath)
+
+    def _rename_sync(self, oldpath: bytes, newpath: bytes) -> None:
+        self._checked_run(f"mv -- {shlex.quote(self._path(oldpath))} {shlex.quote(self._path(newpath))}")
+
+    async def realpath(self, path: bytes) -> bytes:
+        return await self._run_blocking(self._realpath_sync, path)
+
+    def _realpath_sync(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")
+
+    async def setstat(self, path: bytes, attrs: asyncssh.SFTPAttrs) -> None:
+        await self._run_blocking(self._setstat_sync, path, attrs)
+
+    def _setstat_sync(self, path: bytes, attrs: asyncssh.SFTPAttrs) -> None:
+        remote_path = self._path(path)
+        if attrs.size is not None:
+            if attrs.size < 0 or attrs.size > self._max_bytes:
+                raise asyncssh.SFTPFailure(f"file size must be between 0 and {self._max_bytes} bytes")
+            script = "import os,sys; os.truncate(sys.argv[1], int(sys.argv[2]))"
+            self._checked_run(f"python3 -c {shlex.quote(script)} {shlex.quote(remote_path)} {attrs.size}")
+        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_to(self, remote_path: str, temporary_path: Path) -> None:
+        try:
+            result = self._client.download_file(
+                self._backend_session_id(),
+                remote_path,
+                temporary_path,
+                timeout=self._timeout,
+                owner="mock-ssh",
+                max_bytes=self._max_bytes,
+            )
+            if result.bytes_transferred > self._max_bytes:
+                raise HostBridgeError(
+                    "resource_limit",
+                    f"file exceeds mock-ssh limit ({self._max_bytes} bytes)",
+                )
+        except HostBridgeError as exc:
+            raise self._sftp_error(exc) from exc
+
+    def _upload(self, handle: HostBridgeSFTPHandle) -> None:
+        if handle.size > self._max_bytes:
+            raise asyncssh.SFTPFailure(f"file exceeds mock-ssh limit ({self._max_bytes} bytes)")
+        try:
+            self._client.upload_file(
+                self._backend_session_id(),
+                handle.temporary_path,
+                handle.path,
+                mode=handle.mode & 0o7777,
+                timeout=self._timeout,
+                owner="mock-ssh",
+                max_bytes=self._max_bytes,
+            )
+            self._apply_times(handle)
+        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
+
+    def _apply_times(self, handle: HostBridgeSFTPHandle) -> None:
+        if handle.atime is None and handle.mtime is None:
+            return
+        current = self._stat(handle.path, follow=True)
+        atime = int(handle.atime if handle.atime is not None else current.atime or time.time())
+        mtime = int(handle.mtime if handle.mtime is not None else current.mtime or time.time())
+        script = "import os,sys; os.utime(sys.argv[1], (int(sys.argv[2]), int(sys.argv[3])))"
+        self._checked_run(f"python3 -c {shlex.quote(script)} {shlex.quote(handle.path)} {atime} {mtime}")
+
+    def _create_remote(self, remote_path: str, *, mode: int, exclusive: bool) -> None:
+        script = (
+            "import os,sys; "
+            "flags=os.O_WRONLY|os.O_CREAT|(os.O_EXCL if sys.argv[3]=='1' else 0); "
+            "fd=os.open(sys.argv[1],flags,int(sys.argv[2],8)); os.close(fd)"
+        )
+        self._checked_run(
+            f"python3 -c {shlex.quote(script)} {shlex.quote(remote_path)} "
+            f"{shlex.quote(format(mode & 0o7777, 'o'))} {'1' if exclusive else '0'}"
+        )
+
+    def _truncate_remote(self, remote_path: str) -> None:
+        script = "import os,sys; fd=os.open(sys.argv[1],os.O_WRONLY|os.O_TRUNC); os.close(fd)"
+        self._checked_run(f"python3 -c {shlex.quote(script)} {shlex.quote(remote_path)}")
+
+    @staticmethod
+    def _new_temporary_path() -> Path:
+        descriptor, name = tempfile.mkstemp(prefix="hostbridge-sftp-")
+        os.close(descriptor)
+        return Path(name)
+
+    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) -> _TextCommandResult:
+        try:
+            result = self._client.exec(
+                self._backend_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 is None:
+            raise asyncssh.SFTPFailure("remote command completed without an exit status")
+        if result.exit_code != 0:
+            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 _backend_session_id(self) -> str:
+        with self._session_lock:
+            if self._session_id is None:
+                self._session_id = self._session_factory.open()
+            return self._session_id
+
+    @staticmethod
+    def _path(path: bytes) -> str:
+        text = path.decode("utf-8", "surrogateescape")
+        if not text:
+            return "."
+        return posixpath.normpath(text)
+
+    @staticmethod
+    def _handle(file_obj: object) -> HostBridgeSFTPHandle:
+        if not isinstance(file_obj, HostBridgeSFTPHandle):
+            raise asyncssh.SFTPInvalidHandle("invalid HostBridge SFTP handle")
+        return file_obj
+
+    @staticmethod
+    def _sftp_error(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)
+        if "file exists" in lowered or "fileexistserror" in lowered:
+            return asyncssh.SFTPFileAlreadyExists(text)
+        if "directory not empty" in lowered:
+            return asyncssh.SFTPDirNotEmpty(text)
+        if "is a directory" in lowered:
+            return asyncssh.SFTPFileIsADirectory(text)
+        if "not a directory" in lowered:
+            return asyncssh.SFTPNotADirectory(text)
+        return asyncssh.SFTPFailure(text)