@neoline/hostbridge 2.0.3 → 2.0.4

package/src/server_control_mcp/hosts.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-import json
 import os
 import re
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
 
+from .config import load_v1_config
 from .secrets import decrypt_secret
 
 DEFAULT_HOSTS_FILE = Path("~/.hostbridge/hosts.json").expanduser()
@@ -14,6 +14,17 @@ LEGACY_HOSTS_FILE = Path("~/.server_control_mcp/hosts.json").expanduser()
 HOST_ID_PATTERN = re.compile(r"^[A-Za-z0-9_.-]+$")
 
 
+@dataclass(frozen=True, slots=True)
+class HostLimits:
+    max_sessions: int = 16
+    command_timeout_seconds: int = 24 * 60 * 60
+    max_output_bytes: int = 1_000_000
+    max_stdin_bytes: int = 100 * 1024 * 1024
+    max_transfer_bytes: int = 100 * 1024 * 1024
+    max_text_bytes: int = 4 * 1024 * 1024
+    max_shell_bytes: int = 1024 * 1024
+
+
 @dataclass(frozen=True, slots=True)
 class LoginStep:
     expect: str
@@ -21,6 +32,16 @@ class LoginStep:
     send_secret: str | None = None
 
 
+@dataclass(frozen=True, slots=True)
+class SshHostConfig:
+    host: str
+    username: str | None = None
+    port: int = 22
+    identity_file: Path | None = None
+    proxy_jump: str | None = None
+    extra_args: tuple[str, ...] = ()
+
+
 @dataclass(frozen=True, slots=True)
 class HostConfig:
     """A single approved connection target."""
@@ -35,11 +56,23 @@ class HostConfig:
     login_steps: list[LoginStep] = field(default_factory=list)
     ready_expect: str | None = None
     transport: str = "auto"
+    limits: HostLimits = field(default_factory=HostLimits)
+    ssh: SshHostConfig | None = None
 
     @classmethod
-    def script(cls, host_id: str, label: str, script: Path | str, *, transport: str = "auto") -> HostConfig:
+    def script(
+        cls,
+        host_id: str,
+        label: str,
+        script: Path | str,
+        *,
+        transport: str = "auto",
+        limits: HostLimits | None = None,
+    ) -> HostConfig:
         path = Path(script).expanduser()
-        return cls(host_id, label, "/bin/sh", [str(path)], "script", path, transport=transport)
+        return cls(
+            host_id, label, "/bin/sh", [str(path)], "script", path, transport=transport, limits=limits or HostLimits()
+        )
 
     @classmethod
     def command_line(
@@ -54,6 +87,8 @@ class HostConfig:
         login_steps: list[LoginStep] | None = None,
         ready_expect: str | None = None,
         transport: str = "auto",
+        limits: HostLimits | None = None,
+        ssh: SshHostConfig | None = None,
     ) -> HostConfig:
         return cls(
             host_id,
@@ -66,6 +101,8 @@ class HostConfig:
             login_steps or [],
             ready_expect,
             transport,
+            limits or HostLimits(),
+            ssh,
         )
 
     def describe(self) -> dict[str, object]:
@@ -96,25 +133,29 @@ DEFAULT_HOSTS: list[HostConfig] = []
 def load_hosts(config_path: Path | str | None = None) -> list[HostConfig]:
     """Load approved hosts from JSON configuration."""
 
-    configured = config_path or os.environ.get("HOSTBRIDGE_HOSTS_FILE") or os.environ.get("SERVER_CONTROL_MCP_HOSTS_FILE")
-    path = Path(configured).expanduser() if configured else _default_hosts_path()
+    path = configured_hosts_path(config_path)
     if not path.exists():
         return []
 
-    data = json.loads(path.read_text(encoding="utf-8"))
-    raw_hosts = data.get("hosts") if isinstance(data, dict) else None
-    if not isinstance(raw_hosts, list):
-        raise ValueError(f"{path} must contain a top-level 'hosts' list")
+    raw_hosts = load_v1_config(path).hosts
 
-    hosts_by_id: dict[str, HostConfig] = {}
-    order: list[str] = []
+    hosts: list[HostConfig] = []
+    seen: set[str] = set()
     for index, item in enumerate(raw_hosts):
         host = _parse_host(item, path, index)
-        if host.id not in hosts_by_id:
-            order.append(host.id)
-        hosts_by_id[host.id] = host
+        if host.id in seen:
+            raise ValueError(f"{path} contains duplicate host id {host.id!r}")
+        seen.add(host.id)
+        hosts.append(host)
+
+    return hosts
+
 
-    return [hosts_by_id[host_id] for host_id in order]
+def configured_hosts_path(config_path: Path | str | None = None) -> Path:
+    configured = (
+        config_path or os.environ.get("HOSTBRIDGE_HOSTS_FILE") or os.environ.get("SERVER_CONTROL_MCP_HOSTS_FILE")
+    )
+    return Path(configured).expanduser() if configured else _default_hosts_path()
 
 
 def _parse_host(item: object, path: Path, index: int) -> HostConfig:
@@ -137,12 +178,15 @@ def _parse_host(item: object, path: Path, index: int) -> HostConfig:
     transport = str(item.get("transport", "auto")).strip().lower()
     if transport not in {"auto", "pty", "ssh"}:
         raise ValueError(f"{path} hosts[{index}].transport must be one of: auto, pty, ssh")
+    if transport == "ssh" and "ssh" not in item:
+        raise ValueError(f"{path} hosts[{index}].transport ssh requires an ssh connection")
+    limits = _parse_limits(item.get("limits"), path, index)
 
     if "script" in item:
         script = str(item.get("script", "")).strip()
         if not script:
             raise ValueError(f"{path} hosts[{index}].script must not be empty")
-        host = HostConfig.script(host_id, label, script, transport=transport)
+        host = HostConfig.script(host_id, label, script, transport=transport, limits=limits)
         return HostConfig(
             host.id,
             host.label,
@@ -154,19 +198,27 @@ def _parse_host(item: object, path: Path, index: int) -> HostConfig:
             login_steps,
             ready_expect,
             transport,
+            limits,
         )
 
     if "ssh" in item:
+        ssh_config, ssh_args = _parse_ssh(item["ssh"], path, index)
+        if transport == "ssh" and ssh_config.extra_args:
+            raise ValueError(f"{path} hosts[{index}].ssh.extra_args cannot be used with native ssh transport")
+        if transport == "ssh" and login_steps:
+            raise ValueError(f"{path} hosts[{index}].login_steps cannot be used with native ssh transport")
         return HostConfig.command_line(
             host_id,
             label,
             "ssh",
-            _parse_ssh_args(item["ssh"], path, index),
+            ssh_args,
             connection_type="ssh",
             secrets=secrets,
             login_steps=login_steps,
             ready_expect=ready_expect,
             transport=transport,
+            limits=limits,
+            ssh=ssh_config,
         )
 
     command = str(item.get("command", "")).strip()
@@ -184,9 +236,25 @@ def _parse_host(item: object, path: Path, index: int) -> HostConfig:
         login_steps=login_steps,
         ready_expect=ready_expect,
         transport=transport,
+        limits=limits,
     )
 
 
+def _parse_limits(raw_limits: object, path: Path, index: int) -> HostLimits:
+    if raw_limits is None:
+        return HostLimits()
+    if not isinstance(raw_limits, dict):
+        raise ValueError(f"{path} hosts[{index}].limits must be an object")
+    values = {field: getattr(HostLimits(), field) for field in HostLimits.__dataclass_fields__}
+    for name, value in raw_limits.items():
+        if name not in values:
+            raise ValueError(f"{path} hosts[{index}].limits contains unknown limit {name!r}")
+        if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+            raise ValueError(f"{path} hosts[{index}].limits.{name} must be a positive integer")
+        values[name] = value
+    return HostLimits(**values)
+
+
 def _parse_ready_expect(raw_ready_expect: object, path: Path, index: int) -> str | None:
     if raw_ready_expect in (None, ""):
         return None
@@ -218,20 +286,36 @@ def _parse_login_steps(raw_steps: object, path: Path, index: int) -> list[LoginS
         if not expect:
             raise ValueError(f"{path} hosts[{index}].login_steps[{step_index}].expect must not be empty")
         if (send is None) == (send_secret is None):
-            raise ValueError(f"{path} hosts[{index}].login_steps[{step_index}] must define exactly one of send or send_secret")
-        steps.append(LoginStep(expect=expect, send=None if send is None else str(send), send_secret=None if send_secret is None else str(send_secret)))
+            raise ValueError(
+                f"{path} hosts[{index}].login_steps[{step_index}] must define exactly one of send or send_secret"
+            )
+        steps.append(
+            LoginStep(
+                expect=expect,
+                send=None if send is None else str(send),
+                send_secret=None if send_secret is None else str(send_secret),
+            )
+        )
     return steps
 
 
-def _parse_ssh_args(raw_ssh: object, path: Path, index: int) -> list[str]:
+def _parse_ssh(raw_ssh: object, path: Path, index: int) -> tuple[SshHostConfig, list[str]]:
     if not isinstance(raw_ssh, dict):
         raise ValueError(f"{path} hosts[{index}].ssh must be an object")
     target = str(raw_ssh.get("host", "")).strip()
     if not target:
         raise ValueError(f"{path} hosts[{index}].ssh.host must not be empty")
 
+    username: str | None = None
+    host = target
+    if "@" in target:
+        username, host = target.rsplit("@", 1)
+        if not username or not host:
+            raise ValueError(f"{path} hosts[{index}].ssh.host must be [username@]hostname")
+
     args: list[str] = []
     port = raw_ssh.get("port")
+    parsed_port = 22
     if port not in (None, ""):
         if isinstance(port, bool) or not isinstance(port, int | str):
             raise ValueError(f"{path} hosts[{index}].ssh.port must be an integer")
@@ -243,8 +327,9 @@ def _parse_ssh_args(raw_ssh: object, path: Path, index: int) -> list[str]:
             raise ValueError(f"{path} hosts[{index}].ssh.port must be between 1 and 65535")
         args.extend(["-p", str(parsed_port)])
     identity_file = str(raw_ssh.get("identity_file", "")).strip()
+    identity_path = Path(identity_file).expanduser() if identity_file else None
     if identity_file:
-        args.extend(["-i", str(Path(identity_file).expanduser())])
+        args.extend(["-i", str(identity_path)])
     proxy_jump = str(raw_ssh.get("proxy_jump", "")).strip()
     if proxy_jump:
         args.extend(["-J", proxy_jump])
@@ -253,7 +338,17 @@ def _parse_ssh_args(raw_ssh: object, path: Path, index: int) -> list[str]:
         raise ValueError(f"{path} hosts[{index}].ssh.extra_args must be a list of strings")
     args.extend(extra_args)
     args.append(target)
-    return args
+    return (
+        SshHostConfig(
+            host=host,
+            username=username,
+            port=parsed_port,
+            identity_file=identity_path,
+            proxy_jump=proxy_jump or None,
+            extra_args=tuple(extra_args),
+        ),
+        args,
+    )
 
 
 class HostRegistry:
@@ -272,8 +367,7 @@ class HostRegistry:
         except KeyError as exc:
             if not self._hosts:
                 raise KeyError(
-                    "no hosts are configured; create ~/.hostbridge/hosts.json "
-                    "or set HOSTBRIDGE_HOSTS_FILE"
+                    "no hosts are configured; create ~/.hostbridge/hosts.json or set HOSTBRIDGE_HOSTS_FILE"
                 ) from exc
             valid = ", ".join(sorted(self._hosts))
             raise KeyError(f"unknown host '{host_id}'; valid hosts: {valid}") from exc