haproxy-mcp/haproxy_mcp/ssh_ops.py

"""SSH remote execution for HAProxy MCP Server.

When REMOTE_MODE is enabled (SSH_HOST is set), file I/O and subprocess
commands are executed on the remote HAProxy host via SSH.
"""

import subprocess
from .config import (
    SSH_HOST,
    SSH_USER,
    SSH_KEY,
    SSH_PORT,
    REMOTE_MODE,
    SUBPROCESS_TIMEOUT,
    logger,
)


def _ssh_base_cmd() -> list[str]:
    """Build base SSH command with options."""
    cmd = [
        "ssh",
        "-o", "StrictHostKeyChecking=no",
        "-o", "UserKnownHostsFile=/dev/null",
        "-o", "LogLevel=ERROR",
        "-o", "BatchMode=yes",
        "-o", "ConnectTimeout=10",
        "-p", str(SSH_PORT),
    ]
    if SSH_KEY:
        cmd.extend(["-i", SSH_KEY])
    cmd.append(f"{SSH_USER}@{SSH_HOST}")
    return cmd


def remote_exec(command: str, timeout: int = SUBPROCESS_TIMEOUT) -> subprocess.CompletedProcess:
    """Execute a command on the remote host via SSH.

    The command is passed as a single SSH argument so the remote shell
    interprets it as-is (compatible with fish, bash, zsh).

    Args:
        command: Shell command to execute remotely
        timeout: Command timeout in seconds

    Returns:
        CompletedProcess with stdout/stderr

    Raises:
        subprocess.TimeoutExpired: If command times out
        OSError: If SSH command fails to execute
    """
    ssh_cmd = _ssh_base_cmd() + [command]
    logger.debug("SSH exec: %s", command)
    return subprocess.run(
        ssh_cmd,
        capture_output=True,
        text=True,
        timeout=timeout,
    )


def remote_read_file(path: str) -> str:
    """Read a file from the remote host.

    Args:
        path: Absolute file path on remote host

    Returns:
        File contents as string

    Raises:
        FileNotFoundError: If file doesn't exist on remote
        IOError: If read fails
    """
    result = remote_exec(f"cat {path}")
    if result.returncode != 0:
        stderr = result.stderr.strip()
        if "No such file" in stderr:
            raise FileNotFoundError(f"Remote file not found: {path}")
        raise IOError(f"Failed to read remote file {path}: {stderr}")
    return result.stdout


def remote_write_file(path: str, content: str) -> None:
    """Write content to a file on the remote host atomically.

    Uses bash explicitly for bash-specific syntax ($(mktemp), $tmpf).
    The entire script is passed as a single argument to 'bash -c'.

    Args:
        path: Absolute file path on remote host
        content: Content to write

    Raises:
        IOError: If write fails
    """
    ssh_cmd = _ssh_base_cmd()
    # Atomic write via bash (single-quoted to survive remote shell interpretation)
    remote_script = f"tmpf=$(mktemp {path}.tmp.XXXXXX) && cat >\"$tmpf\" && mv \"$tmpf\" {path}"
    ssh_cmd.append(f"bash -c '{remote_script}'")

    logger.debug("SSH write: %s (%d bytes)", path, len(content))
    result = subprocess.run(
        ssh_cmd,
        input=content,
        capture_output=True,
        text=True,
        timeout=SUBPROCESS_TIMEOUT,
    )
    if result.returncode != 0:
        raise IOError(f"Failed to write remote file {path}: {result.stderr.strip()}")


def remote_file_exists(path: str) -> bool:
    """Check if a file exists on the remote host.

    Args:
        path: Absolute file path on remote host

    Returns:
        True if file exists
    """
    result = remote_exec(f"test -f {path} && echo yes || echo no")
    return result.stdout.strip() == "yes"


def _scp_base_cmd() -> list[str]:
    """Build base SCP command with options."""
    cmd = [
        "scp",
        "-o", "StrictHostKeyChecking=no",
        "-o", "UserKnownHostsFile=/dev/null",
        "-o", "LogLevel=ERROR",
        "-o", "BatchMode=yes",
        "-o", "ConnectTimeout=10",
        "-P", str(SSH_PORT),
    ]
    if SSH_KEY:
        cmd.extend(["-i", SSH_KEY])
    return cmd


def remote_download_file(remote_path: str, local_path: str) -> bool:
    """Download a binary file from the remote host via SCP.

    Args:
        remote_path: Absolute file path on remote host
        local_path: Absolute local file path to write to

    Returns:
        True if downloaded successfully, False if file doesn't exist
    """
    cmd = _scp_base_cmd() + [f"{SSH_USER}@{SSH_HOST}:{remote_path}", local_path]
    logger.debug("SCP download: %s -> %s", remote_path, local_path)
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=SUBPROCESS_TIMEOUT)
    if result.returncode != 0:
        logger.debug("SCP download failed: %s", result.stderr.strip())
        return False
    return True


def remote_upload_file(local_path: str, remote_path: str) -> None:
    """Upload a binary file to the remote host via SCP.

    Args:
        local_path: Absolute local file path to upload
        remote_path: Absolute file path on remote host

    Raises:
        IOError: If upload fails
    """
    cmd = _scp_base_cmd() + [local_path, f"{SSH_USER}@{SSH_HOST}:{remote_path}"]
    logger.debug("SCP upload: %s -> %s", local_path, remote_path)
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=SUBPROCESS_TIMEOUT)
    if result.returncode != 0:
        raise IOError(f"SCP upload failed: {result.stderr.strip()}")


def run_command(args: list[str], timeout: int = SUBPROCESS_TIMEOUT) -> subprocess.CompletedProcess:
    """Execute a command locally or remotely based on REMOTE_MODE.

    Args:
        args: Command and arguments as list
        timeout: Command timeout in seconds

    Returns:
        CompletedProcess with stdout/stderr
    """
    if REMOTE_MODE:
        # Join args into a shell command for SSH
        # Quote arguments that contain spaces
        quoted = []
        for a in args:
            if " " in a or "'" in a or '"' in a:
                quoted.append(f"'{a}'")
            else:
                quoted.append(a)
        return remote_exec(" ".join(quoted), timeout=timeout)
    else:
        return subprocess.run(
            args,
            capture_output=True,
            text=True,
            timeout=timeout,
        )