hermes-agent/hermes_cli/ssh_tunnel.py

"""Managed SSH tunnels for configured OpenAI-compatible endpoints.

A tunnel is process-local and reused by every agent in the same Hermes process.
The persisted provider URL remains the address as seen from the SSH host; only
the runtime URL is rewritten to a loopback port selected by the kernel.
"""

from __future__ import annotations

import atexit
import hashlib
import logging
import os
import shutil
import socket
import subprocess
import threading
import time
from dataclasses import dataclass
from typing import Any
from urllib.parse import urlparse, urlunparse

logger = logging.getLogger(__name__)

_CONNECT_TIMEOUT_SECONDS = 15
_START_TIMEOUT_SECONDS = 5
_TUNNELS: dict[str, "ManagedSshTunnel"] = {}
_TUNNELS_LOCK = threading.Lock()
_CONTROL_CHARS = frozenset(chr(i) for i in (*range(32), 127))


def _has_control_chars(value: str) -> bool:
    return any(char in _CONTROL_CHARS for char in value)


def _validate_ssh_value(name: str, value: str, *, required: bool = False) -> str:
    clean = str(value or "").strip()
    if required and not clean:
        raise ValueError(f"SSH tunnel requires {name}.")
    if clean and (_has_control_chars(clean) or clean.startswith("-")):
        raise ValueError(f"Unsafe SSH tunnel {name}.")
    return clean


@dataclass(frozen=True)
class SshTunnelConfig:
    host: str
    user: str = ""
    port: int = 22
    key_path: str = ""

    @classmethod
    def from_dict(cls, raw: Any) -> "SshTunnelConfig | None":
        if not isinstance(raw, dict):
            return None
        host = _validate_ssh_value("host", raw.get("host", ""), required=True)
        user = _validate_ssh_value("user", raw.get("user", ""))
        if "@" in host and not user:
            user, host = host.split("@", 1)
            user = _validate_ssh_value("user", user, required=True)
            host = _validate_ssh_value("host", host, required=True)
        elif "@" in host:
            raise ValueError("SSH tunnel host must not include a user when SSH user is set separately.")
        key_path = _validate_ssh_value("key path", raw.get("key_path", ""))
        try:
            port = int(raw.get("port") or 22)
        except (TypeError, ValueError) as exc:
            raise ValueError("SSH tunnel port must be an integer between 1 and 65535.") from exc
        if not 1 <= port <= 65535:
            raise ValueError("SSH tunnel port must be between 1 and 65535.")
        return cls(host=host, user=user, port=port, key_path=key_path)


def _pick_local_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("127.0.0.1", 0))
        return int(sock.getsockname()[1])


def _endpoint_target(base_url: str) -> tuple[str, int, str]:
    parsed = urlparse(str(base_url or "").strip())
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise ValueError("SSH-tunneled endpoint URL must include http(s) scheme and host.")
    try:
        remote_port = parsed.port or (443 if parsed.scheme == "https" else 80)
    except ValueError as exc:
        raise ValueError("SSH-tunneled endpoint URL has an invalid port.") from exc
    return parsed.hostname, remote_port, parsed.geturl().rstrip("/")


class ManagedSshTunnel:
    def __init__(self, config: SshTunnelConfig, remote_host: str, remote_port: int):
        self.config = config
        self.remote_host = remote_host
        self.remote_port = remote_port
        self.local_port: int | None = None
        self.process: subprocess.Popen[bytes] | None = None

    @property
    def target(self) -> str:
        return f"{self.config.user + '@' if self.config.user else ''}{self.config.host}"

    def _args(self, local_port: int) -> list[str]:
        args = [
            "ssh",
            "-N",
            "-o", "BatchMode=yes",
            "-o", "StrictHostKeyChecking=accept-new",
            "-o", "ExitOnForwardFailure=yes",
            "-o", f"ConnectTimeout={_CONNECT_TIMEOUT_SECONDS}",
        ]
        if self.config.port != 22:
            args.extend(["-p", str(self.config.port)])
        if self.config.key_path:
            args.extend(["-i", self.config.key_path])
        args.extend(["-L", f"127.0.0.1:{local_port}:{self.remote_host}:{self.remote_port}", "--", self.target])
        return args

    def start(self) -> int:
        if self.local_port and self.process and self.process.poll() is None:
            return self.local_port
        if not shutil.which("ssh"):
            raise RuntimeError("SSH is not installed or not in PATH. Install an OpenSSH client first.")
        if self.config.key_path and not os.path.isfile(os.path.expanduser(self.config.key_path)):
            raise RuntimeError(f"SSH identity file does not exist: {self.config.key_path}")

        for _ in range(3):
            local_port = _pick_local_port()
            process = subprocess.Popen(
                self._args(local_port),
                stdin=subprocess.DEVNULL,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
            )
            deadline = time.monotonic() + _START_TIMEOUT_SECONDS
            while time.monotonic() < deadline:
                if process.poll() is not None:
                    process.communicate()
                    raise RuntimeError("SSH tunnel failed to start. Check the SSH connection and credentials.")
                try:
                    with socket.create_connection(("127.0.0.1", local_port), timeout=0.15):
                        self.process = process
                        self.local_port = local_port
                        logger.info("SSH tunnel ready on loopback port %s", local_port)
                        return local_port
                except OSError:
                    time.sleep(0.05)
            process.terminate()
            try:
                process.wait(timeout=1)
            except subprocess.TimeoutExpired:
                process.kill()
        raise RuntimeError("SSH tunnel could not bind an automatically assigned local port.")

    def close(self) -> None:
        if self.process and self.process.poll() is None:
            self.process.terminate()
            try:
                self.process.wait(timeout=3)
            except subprocess.TimeoutExpired:
                self.process.kill()
        self.process = None
        self.local_port = None


def resolve_ssh_tunnel_url(base_url: str, ssh_tunnel: Any) -> str:
    """Start/reuse a configured tunnel and return its loopback runtime URL."""
    config = SshTunnelConfig.from_dict(ssh_tunnel)
    if config is None:
        return base_url.rstrip("/")
    remote_host, remote_port, normalized_url = _endpoint_target(base_url)
    identity = hashlib.sha256(repr((config, remote_host, remote_port)).encode()).hexdigest()
    with _TUNNELS_LOCK:
        tunnel = _TUNNELS.get(identity)
        if tunnel is None:
            tunnel = ManagedSshTunnel(config, remote_host, remote_port)
            _TUNNELS[identity] = tunnel
        local_port = tunnel.start()
    parsed = urlparse(normalized_url)
    return urlunparse(parsed._replace(netloc=f"127.0.0.1:{local_port}")).rstrip("/")


def close_ssh_tunnels() -> None:
    with _TUNNELS_LOCK:
        tunnels = list(_TUNNELS.values())
        _TUNNELS.clear()
    for tunnel in tunnels:
        tunnel.close()


atexit.register(close_ssh_tunnels)