security-alert-center/backend/tests/test_ssh_connect.py

"""SSH connect service tests (paramiko mocked via sys.modules)."""

import sys
from unittest.mock import MagicMock

from app.models import Host
from app.services import ssh_connect
from app.services.ssh_connect import (
    HostNotLinuxError,
    HostTargetMissingError,
    _remote_shell_command,
    iter_ssh_targets,
    run_ssh_command,
    run_ssh_monitor_update,
)


class _AuthError(Exception):
    pass


def _install_fake_paramiko(monkeypatch, *, connect_side_effect=None, exec_setup=None):
    mock_client_cls = MagicMock()
    client = MagicMock()
    if connect_side_effect is not None:
        client.connect.side_effect = connect_side_effect
    if exec_setup is not None:
        exec_setup(client)
    mock_client_cls.return_value = client

    fake = MagicMock()
    fake.SSHClient = mock_client_cls
    fake.AutoAddPolicy = MagicMock()
    fake.AuthenticationException = _AuthError
    monkeypatch.setitem(sys.modules, "paramiko", fake)
    return client


def test_remote_shell_command_non_root_probe_has_no_sudo():
    cmd = _remote_shell_command("deploy", "hostname", "secret", need_root=False)
    assert "sudo" not in cmd
    assert "hostname" in cmd


def test_remote_shell_command_non_root_privileged_uses_sudo_s():
    cmd = _remote_shell_command("deploy", "/opt/scripts/update_ssh_monitor.sh", "secret", need_root=True)
    assert "sudo -S" in cmd
    assert "secret" in cmd
    assert "update_ssh_monitor.sh" in cmd


def test_remote_shell_command_root_skips_sudo_even_when_need_root():
    cmd = _remote_shell_command("root", "/opt/scripts/update_ssh_monitor.sh", "secret", need_root=True)
    assert "sudo" not in cmd


def test_probe_ssh_connection_non_root_does_not_use_sudo(monkeypatch):
    captured: list[str] = []

    def setup(client):
        stdout = MagicMock()
        stdout.read.return_value = b"ubabuba\n"
        stdout.channel.recv_exit_status.return_value = 0
        stderr = MagicMock()
        stderr.read.return_value = b""

        def capture_exec(cmd, **kwargs):
            captured.append(cmd)
            return (None, stdout, stderr)

        client.exec_command.side_effect = capture_exec

    _install_fake_paramiko(monkeypatch, exec_setup=setup)

    result = ssh_connect.test_ssh_connection(target="10.10.36.9", user="deploy", password="pw")
    assert result.ok is True
    assert captured
    assert "sudo" not in captured[0]


def test_iter_ssh_targets_skips_human_display_name(monkeypatch):
    monkeypatch.setattr(ssh_connect, "_ssh_name_resolves", lambda name: True)
    host = Host(
        hostname="ubabuba",
        display_name="Ubabuba Kalina (10.10.36.9)",
        os_family="linux",
        product="ssh-monitor",
        ipv4="10.10.36.9",
    )
    targets = iter_ssh_targets(host)
    assert "Ubabuba Kalina (10.10.36.9)" not in targets
    assert "10.10.36.9" in targets


def test_iter_ssh_targets_skips_unresolvable_hostname(monkeypatch):
    def resolves(name: str) -> bool:
        return ssh_connect._is_ipv4(name)

    monkeypatch.setattr(ssh_connect, "_ssh_name_resolves", resolves)
    host = Host(
        hostname="ubabuba",
        os_family="linux",
        product="ssh-monitor",
        ipv4="10.10.36.9",
    )
    targets = iter_ssh_targets(host)
    assert targets == ["10.10.36.9"]


def test_iter_ssh_targets_rejects_windows():
    host = Host(hostname="PC", os_family="windows", product="rdp-login-monitor", ipv4="1.2.3.4")
    try:
        iter_ssh_targets(host)
        raise AssertionError("expected HostNotLinuxError")
    except HostNotLinuxError:
        pass


def test_iter_ssh_targets_requires_address():
    host = Host(hostname="", os_family="linux", product="ssh-monitor", ipv4=None)
    try:
        iter_ssh_targets(host)
        raise AssertionError("expected HostTargetMissingError")
    except HostTargetMissingError:
        pass


def test_probe_ssh_connection_success(monkeypatch):
    def setup(client):
        stdout = MagicMock()
        stdout.read.return_value = b"ubabuba\n"
        stdout.channel.recv_exit_status.return_value = 0
        stderr = MagicMock()
        stderr.read.return_value = b""
        client.exec_command.return_value = (None, stdout, stderr)

    client = _install_fake_paramiko(monkeypatch, exec_setup=setup)

    result = ssh_connect.test_ssh_connection(target="ubabuba", user="root", password="pw")
    assert result.ok is True
    assert "hostname=ubabuba" in result.message
    client.connect.assert_called_once()
    client.close.assert_called_once()


def test_run_ssh_command_auth_failure(monkeypatch):
    client = _install_fake_paramiko(monkeypatch, connect_side_effect=_AuthError("bad creds"))

    result = run_ssh_command(target="10.0.0.1", user="root", password="wrong", remote_cmd="hostname")
    assert result.ok is False
    assert "auth failed" in result.message.lower()
    client.close.assert_called_once()


def test_run_ssh_monitor_update_missing_script(monkeypatch):
    monkeypatch.setattr(
        ssh_connect,
        "run_ssh_command",
        lambda **kwargs: ssh_connect.SshCommandResult(
            ok=True,
            message="ok",
            target="h",
            stdout="missing\n",
        ),
    )

    result = run_ssh_monitor_update(target="ubabuba", user="root", password="pw")
    assert result.ok is False
    assert "not found" in result.message.lower()


def test_run_ssh_monitor_update_runs_script(monkeypatch):
    calls = 0

    def fake_run(**kwargs):
        nonlocal calls
        calls += 1
        if calls == 1:
            return ssh_connect.SshCommandResult(
                ok=True,
                message="ok",
                target="h",
                stdout="ready\n",
            )
        if calls == 2:
            assert kwargs["remote_cmd"] == "/opt/scripts/update_ssh_monitor.sh"
            assert kwargs.get("need_root") is True
            return ssh_connect.SshCommandResult(
                ok=True,
                message="SSH OK",
                target="ubabuba",
                stdout="SUMMARY updated 2.1.0-SAC\n",
                exit_code=0,
            )
        return ssh_connect.SshCommandResult(
            ok=True,
            message="ok",
            target="ubabuba",
            stdout='2.1.0-SAC\n',
            exit_code=0,
        )

    monkeypatch.setattr(ssh_connect, "run_ssh_command", fake_run)

    result = run_ssh_monitor_update(target="ubabuba", user="root", password="pw")
    assert result.ok is True
    assert result.agent_version == "2.1.0-SAC"
    assert calls == 3


def test_parse_ssh_monitor_version_text():
    assert ssh_connect.parse_ssh_monitor_version_text("SSH_MONITOR_VERSION=2.1.0-SAC") == "2.1.0-SAC"
    assert ssh_connect.parse_ssh_monitor_version_text("SUMMARY updated 2.1.1-SAC done") == "2.1.1-SAC"
    assert ssh_connect.parse_ssh_monitor_version_text("no version") is None