security-alert-center/backend/app/services/ssh_connect.py

"""SSH connectivity and remote commands for Linux hosts."""

from __future__ import annotations

import socket
from dataclasses import dataclass

from app.models import Host

SSH_MONITOR_UPDATE_SCRIPT = "/opt/scripts/update_ssh_monitor.sh"
SSH_OUTPUT_MAX_LEN = 12_000


class LinuxAdminNotConfiguredError(Exception):
    pass


class HostNotLinuxError(Exception):
    pass


class HostTargetMissingError(Exception):
    pass


@dataclass(frozen=True)
class SshCommandResult:
    ok: bool
    message: str
    target: str
    stdout: str = ""
    stderr: str = ""
    exit_code: int | None = None


def is_linux_host(host: Host) -> bool:
    if (host.os_family or "").strip().lower() == "linux":
        return True
    return (host.product or "").strip() == "ssh-monitor"


def iter_ssh_targets(host: Host) -> list[str]:
    if not is_linux_host(host):
        raise HostNotLinuxError("Host is not Linux")

    seen: set[str] = set()
    ordered: list[str] = []

    def add(value: str | None) -> None:
        text = (value or "").strip()
        if not text or text.casefold() in seen:
            return
        seen.add(text.casefold())
        ordered.append(text)

    add(host.hostname)
    add(host.display_name)
    add(host.ipv4)
    if not ordered:
        raise HostTargetMissingError("Host has no hostname or IPv4 for SSH")
    return ordered


def _truncate_output(text: str) -> str:
    if len(text) <= SSH_OUTPUT_MAX_LEN:
        return text
    return text[:SSH_OUTPUT_MAX_LEN] + "\n… (truncated)"


def _shell_single_quote(value: str) -> str:
    return "'" + value.replace("'", "'\"'\"'") + "'"


def _remote_shell_command(
    user: str,
    remote_cmd: str,
    password: str,
    *,
    need_root: bool = False,
) -> str:
    """Build remote shell command. Sudo only when need_root and user is not root."""
    safe_cmd = _shell_single_quote(remote_cmd)
    if user == "root" or not need_root:
        return f"bash -lc {safe_cmd}"

    safe_pw = _shell_single_quote(password)
    inner = f"printf '%s\\n' {safe_pw} | sudo -S -p '' bash -lc {safe_cmd}"
    return f"bash -lc {_shell_single_quote(inner)}"


def run_ssh_command(
    *,
    target: str,
    user: str,
    password: str,
    remote_cmd: str,
    connect_timeout_sec: int = 15,
    command_timeout_sec: int = 600,
    need_root: bool = False,
) -> SshCommandResult:
    target = target.strip()
    user = user.strip()
    if not target or not user or not password:
        raise LinuxAdminNotConfiguredError("Linux SSH admin credentials or target missing")

    try:
        import paramiko
    except ImportError as exc:
        raise RuntimeError("paramiko is not installed on SAC server") from exc

    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    shell_cmd = _remote_shell_command(user, remote_cmd, password, need_root=need_root)

    try:
        client.connect(
            hostname=target,
            username=user,
            password=password,
            timeout=connect_timeout_sec,
            allow_agent=False,
            look_for_keys=False,
        )
        _stdin, stdout, stderr = client.exec_command(shell_cmd, timeout=command_timeout_sec)
        exit_code = stdout.channel.recv_exit_status()
        out_text = _truncate_output(stdout.read().decode("utf-8", errors="replace"))
        err_text = _truncate_output(stderr.read().decode("utf-8", errors="replace"))
    except paramiko.AuthenticationException:
        return SshCommandResult(
            ok=False,
            message=f"SSH auth failed for {user}@{target}",
            target=target,
        )
    except (socket.timeout, TimeoutError):
        return SshCommandResult(
            ok=False,
            message=f"SSH timeout ({connect_timeout_sec}s connect / {command_timeout_sec}s cmd) to {target}",
            target=target,
        )
    except Exception as exc:
        return SshCommandResult(
            ok=False,
            message=f"SSH error ({target}): {exc}",
            target=target,
        )
    finally:
        client.close()

    ok = exit_code == 0
    if ok:
        message = f"SSH OK ({target}), exit 0"
        if out_text.strip():
            message = f"{message}\n{out_text.strip().splitlines()[0][:200]}"
    else:
        detail = err_text.strip() or out_text.strip() or f"exit code {exit_code}"
        message = f"SSH command failed ({target}): {detail[:500]}"

    return SshCommandResult(
        ok=ok,
        message=message,
        target=target,
        stdout=out_text,
        stderr=err_text,
        exit_code=exit_code,
    )


def test_ssh_connection(
    *,
    target: str,
    user: str,
    password: str,
    timeout_sec: int = 15,
) -> SshCommandResult:
    result = run_ssh_command(
        target=target,
        user=user,
        password=password,
        remote_cmd="hostname",
        connect_timeout_sec=timeout_sec,
        command_timeout_sec=timeout_sec,
    )
    if result.ok and result.stdout.strip():
        host = result.stdout.strip().splitlines()[0]
        return SshCommandResult(
            ok=True,
            message=f"SSH OK, hostname={host}",
            target=target,
            stdout=result.stdout,
            stderr=result.stderr,
            exit_code=result.exit_code,
        )
    return result


def run_ssh_monitor_update(
    *,
    target: str,
    user: str,
    password: str,
) -> SshCommandResult:
    script = SSH_MONITOR_UPDATE_SCRIPT
    check_cmd = f"test -x {script} && echo ready || echo missing"
    probe = run_ssh_command(
        target=target,
        user=user,
        password=password,
        remote_cmd=check_cmd,
        command_timeout_sec=30,
    )
    if not probe.ok or "missing" in probe.stdout:
        return SshCommandResult(
            ok=False,
            message=f"Update script not found or not executable: {script} on {target}",
            target=target,
            stdout=probe.stdout,
            stderr=probe.stderr,
            exit_code=probe.exit_code,
        )

    return run_ssh_command(
        target=target,
        user=user,
        password=password,
        remote_cmd=script,
        command_timeout_sec=900,
        need_root=True,
    )