@neoline/hostbridge 2.0.3 → 2.0.5

package/src/server_control_mcp/secrets.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import os
-from contextlib import suppress
 from pathlib import Path
 from typing import Any
 
 from cryptography.fernet import Fernet, InvalidToken
 
+from .secure_files import create_private_file, validate_private_file
+
 DEFAULT_KEY_FILE = Path("~/.hostbridge/key").expanduser()
 LEGACY_KEY_FILE = Path("~/.server_control_mcp/key").expanduser()
 SCHEME = "fernet.v1"
@@ -24,13 +25,12 @@ def key_file_path() -> Path:
 def _load_or_create_key() -> bytes:
     path = key_file_path()
     if path.exists():
+        validate_private_file(path)
         return path.read_bytes().strip()
-    path.parent.mkdir(parents=True, exist_ok=True)
     key = Fernet.generate_key()
-    path.write_bytes(key + b"\n")
-    with suppress(OSError):
-        path.chmod(0o600)
-    return key
+    if create_private_file(path, key + b"\n"):
+        return key
+    return path.read_bytes().strip()
 
 
 def encrypt_secret(value: str) -> dict[str, str]:

package/src/server_control_mcp/secure_files.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import errno
+import os
+import stat
+import tempfile
+from contextlib import suppress
+from pathlib import Path
+
+
+def ensure_private_directory(path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with suppress(FileExistsError):
+        os.mkdir(path, 0o700)
+    metadata = path.lstat()
+    if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode):
+        raise RuntimeError(f"secure directory is not a real directory: {path}")
+    if metadata.st_uid != os.getuid():
+        raise PermissionError(f"secure directory is not owned by the current user: {path}")
+    path.chmod(0o700)
+
+
+def validate_private_file(path: Path) -> None:
+    metadata = path.lstat()
+    if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISREG(metadata.st_mode):
+        raise RuntimeError(f"private file is not a regular file: {path}")
+    if metadata.st_uid != os.getuid():
+        raise PermissionError(f"private file is not owned by the current user: {path}")
+    path.chmod(0o600)
+
+
+def read_private_text(path: Path) -> str:
+    nofollow = getattr(os, "O_NOFOLLOW", None)
+    if nofollow is None:
+        raise RuntimeError("secure private file reads require O_NOFOLLOW")
+    flags = os.O_RDONLY | nofollow
+    if hasattr(os, "O_CLOEXEC"):
+        flags |= os.O_CLOEXEC
+    try:
+        descriptor = os.open(path, flags)
+    except OSError as exc:
+        if exc.errno == errno.ELOOP:
+            raise RuntimeError(f"private file is not a regular file: {path}") from exc
+        raise
+    try:
+        metadata = os.fstat(descriptor)
+        if not stat.S_ISREG(metadata.st_mode):
+            raise RuntimeError(f"private file is not a regular file: {path}")
+        if metadata.st_uid != os.getuid():
+            raise PermissionError(f"private file is not owned by the current user: {path}")
+        os.fchmod(descriptor, 0o600)
+        with os.fdopen(descriptor, "r", encoding="utf-8") as handle:
+            descriptor = -1
+            return handle.read()
+    finally:
+        if descriptor >= 0:
+            os.close(descriptor)
+
+
+def create_private_file(path: Path, data: bytes) -> bool:
+    ensure_private_directory(path.parent)
+    descriptor, name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
+    temporary = Path(name)
+    try:
+        os.fchmod(descriptor, 0o600)
+        _write_all(descriptor, data)
+        os.fsync(descriptor)
+        os.close(descriptor)
+        descriptor = -1
+        try:
+            os.link(temporary, path)
+        except FileExistsError:
+            validate_private_file(path)
+            return False
+        validate_private_file(path)
+        return True
+    finally:
+        if descriptor >= 0:
+            os.close(descriptor)
+        temporary.unlink(missing_ok=True)
+
+
+def replace_private_file(path: Path, data: bytes) -> None:
+    ensure_private_directory(path.parent)
+    descriptor, name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
+    temporary = Path(name)
+    try:
+        os.fchmod(descriptor, 0o600)
+        _write_all(descriptor, data)
+        os.fsync(descriptor)
+        os.close(descriptor)
+        descriptor = -1
+        temporary.replace(path)
+        validate_private_file(path)
+    finally:
+        if descriptor >= 0:
+            os.close(descriptor)
+        temporary.unlink(missing_ok=True)
+
+
+def open_private_append(path: Path):
+    ensure_private_directory(path.parent)
+    flags = os.O_WRONLY | os.O_APPEND | os.O_CREAT
+    if hasattr(os, "O_NOFOLLOW"):
+        flags |= os.O_NOFOLLOW
+    descriptor = os.open(path, flags, 0o600)
+    metadata = os.fstat(descriptor)
+    if not stat.S_ISREG(metadata.st_mode) or metadata.st_uid != os.getuid():
+        os.close(descriptor)
+        raise PermissionError(f"private log is not owned by the current user: {path}")
+    os.fchmod(descriptor, 0o600)
+    return os.fdopen(descriptor, "a", encoding="utf-8")
+
+
+def _write_all(descriptor: int, data: bytes) -> None:
+    remaining = memoryview(data)
+    while remaining:
+        written = os.write(descriptor, remaining)
+        if written <= 0:
+            raise OSError("private file write made no progress")
+        remaining = remaining[written:]

package/src/server_control_mcp/server.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
+import asyncio
 import os
 import sys
-from collections.abc import Callable
+import threading
+from collections.abc import AsyncIterator, Callable
+from contextlib import asynccontextmanager
 from dataclasses import asdict
+from pathlib import Path
 from typing import Any, TypeVar
 
 from mcp.server.fastmcp import FastMCP
@@ -16,12 +20,39 @@ DEFAULT_CONNECT_TIMEOUT = 30
 DEFAULT_MAX_TRANSFER_BYTES = 100 * 1024 * 1024
 MAX_TRANSFER_BYTES = 1024 * 1024 * 1024
 
-mcp = FastMCP("hostbridge")
+_client_instance: HostBridgeClient | None = None
+_client_lock = threading.Lock()
+
+
+@asynccontextmanager
+async def _mcp_lifespan(_server: FastMCP[Any]) -> AsyncIterator[None]:
+    try:
+        yield None
+    finally:
+        await asyncio.to_thread(_close_client)
+
+
+mcp = FastMCP("hostbridge", lifespan=_mcp_lifespan)
 T = TypeVar("T")
 
 
 def _client() -> HostBridgeClient:
-    return HostBridgeClient()
+    global _client_instance
+    if daemon.start_background(quiet=True) != 0:
+        raise RuntimeError(f"hostbridge daemon failed to start; see {daemon.log_file()}")
+    with _client_lock:
+        if _client_instance is None or _client_instance.closed:
+            _client_instance = HostBridgeClient()
+        return _client_instance
+
+
+def _close_client() -> None:
+    global _client_instance
+    with _client_lock:
+        client = _client_instance
+        _client_instance = None
+    if client is not None:
+        client.close()
 
 
 def _owner() -> str | None:
@@ -48,14 +79,7 @@ def _request(
     owner = _owner()
     if owner is not None:
         payload.setdefault("owner", owner)
-    try:
-        return _client().request(method, payload, timeout=timeout)
-    except HostBridgeError as exc:
-        if exc.code != "connection_lost" or os.environ.get("HOSTBRIDGE_NPM_LAUNCHER"):
-            raise
-        if daemon.start_background(quiet=True) != 0:
-            raise RuntimeError(f"hostbridge daemon failed to start; see {daemon.log_file()}") from exc
-        return _client().request(method, payload, timeout=timeout)
+    return _client().request(method, payload, timeout=timeout)
 
 
 def _mode(value: str | int) -> int:
@@ -78,9 +102,9 @@ def hosts_list() -> dict[str, object]:
 @mcp.tool()
 def session_open(host_id: str, connect_timeout: int = DEFAULT_CONNECT_TIMEOUT) -> dict[str, object]:
     """Open a persistent session for an allowlisted host."""
-    _ = max(1, min(int(connect_timeout), 180))
+    safe_timeout = max(1, min(int(connect_timeout), 180))
     return _call(
-        lambda: _client().open_session(host_id, owner=_owner()),
+        lambda: _client().open_session(host_id, owner=_owner(), connect_timeout=safe_timeout),
         lambda value: {"ok": True, "session": asdict(value)},
     )
 
@@ -92,10 +116,10 @@ def sessions_list() -> dict[str, object]:
 
 
 @mcp.tool()
-def session_close(session_id: str, force: bool = False) -> dict[str, object]:
+def session_close(session_id: str) -> dict[str, object]:
     """Close a session."""
     return _call(
-        lambda: _client().close_session(session_id, owner=_owner(), force=bool(force)),
+        lambda: _client().close_session(session_id, owner=_owner()),
         lambda _: {"ok": True, "session_id": session_id, "closed": True},
     )
 
@@ -204,9 +228,18 @@ def file_upload(
     session_id: str, local_path: str, remote_path: str, mode: str = "0644", max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES
 ) -> dict[str, object]:
     """Stream a local file to the remote host."""
-    del max_bytes
+    safe_limit = _limit(max_bytes)
+    source = Path(local_path)
+    try:
+        source_size = source.stat().st_size
+    except OSError as exc:
+        return _error(exc)
+    if source_size > safe_limit:
+        return _error(HostBridgeError("resource_limit", f"upload exceeds {safe_limit} bytes"))
     return _call(
-        lambda: _client().upload_file(session_id, local_path, remote_path, mode=_mode(mode), owner=_owner()),
+        lambda: _client().upload_file(
+            session_id, source, remote_path, mode=_mode(mode), owner=_owner(), max_bytes=safe_limit
+        ),
         lambda value: {"ok": True, "transfer": asdict(value)},
     )
 
@@ -216,16 +249,16 @@ def file_download(
     session_id: str, remote_path: str, local_path: str, max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES
 ) -> dict[str, object]:
     """Stream a remote file to a local path."""
-    del max_bytes
+    safe_limit = _limit(max_bytes)
     return _call(
-        lambda: _client().download_file(session_id, remote_path, local_path, owner=_owner()),
+        lambda: _client().download_file(session_id, remote_path, local_path, owner=_owner(), max_bytes=safe_limit),
         lambda value: {"ok": True, "transfer": asdict(value)},
     )
 
 
 def main(argv: list[str] | None = None) -> Any:
     args = sys.argv[1:] if argv is None else argv
-    if args and (args[0] in cli.SETUP_COMMANDS or args[0] in {"-h", "--help"}):
+    if args:
         return cli.main(args)
     return mcp.run(transport="stdio")