ClassicWasTaken--HaloSteamFrameMod/tests/test_ssh.py

import threading
from unittest.mock import Mock

import paramiko
import pytest

from halo_frame_installer.ssh import Settings, SSHConnection, SSHError, _VerifyHost, fingerprint, validate_host


@pytest.mark.parametrize("host", ["192.168.1.12", "frame.local", "frame", "::1"])
def test_host_accepts_network_addresses(host):
    assert validate_host(host) == host


@pytest.mark.parametrize("host", ["", "ssh://frame", "frame;rm -rf /", "frame\nother", "-host", "user@frame", "a..b"])
def test_host_rejects_command_or_url_input(host):
    with pytest.raises(ValueError):
        validate_host(host)


def test_settings_does_not_display_password():
    value = Settings("frame", "my-private-password")
    assert "my-private-password" not in repr(value)
    with pytest.raises(ValueError):
        Settings("frame", "password", username="root")


def test_new_host_requires_approval_before_trust():
    key = Mock()
    key.asbytes.return_value = b"public key"
    key.get_name.return_value = "ssh-ed25519"
    client = Mock()
    policy = _VerifyHost(Settings("frame", "password", accept_host_key=lambda host, kind, fp: False))
    with pytest.raises(SSHError, match="not approved"):
        policy.missing_host_key(client, "frame", key)
    client.get_host_keys.assert_not_called()


def test_matching_saved_host_key_works_without_another_prompt():
    key = Mock()
    key.asbytes.return_value = b"public key"
    key.get_name.return_value = "ssh-ed25519"
    prompt = Mock(side_effect=AssertionError("must not prompt"))
    policy = _VerifyHost(Settings("frame", "password", known_host_fingerprint=fingerprint(key), accept_host_key=prompt))
    client = Mock()
    policy.missing_host_key(client, "frame", key)
    assert policy.accepted_fingerprint == fingerprint(key)
    prompt.assert_not_called()


def test_changed_host_key_cannot_be_approved_over_a_saved_key():
    key = Mock()
    key.asbytes.return_value = b"different public key"
    prompt = Mock(return_value=True)
    policy = _VerifyHost(Settings("frame", "password", known_host_fingerprint="SHA256:old", accept_host_key=prompt))
    with pytest.raises(SSHError, match="changed"):
        policy.missing_host_key(Mock(), "frame", key)
    prompt.assert_not_called()


def test_connect_uses_memory_password_no_agent_or_disk_keys(monkeypatch):
    client = Mock()
    monkeypatch.setattr(paramiko, "SSHClient", lambda: client)
    connection = SSHConnection(Settings("frame", "private"))
    connection.connect()
    kwargs = client.connect.call_args.kwargs
    assert kwargs["password"] == "private"
    assert kwargs["allow_agent"] is False and kwargs["look_for_keys"] is False
    client.load_system_host_keys.assert_not_called()
    client.save_host_keys.assert_not_called()


def test_authentication_error_does_not_include_secret(monkeypatch):
    client = Mock()
    client.connect.side_effect = paramiko.AuthenticationException("private password text")
    monkeypatch.setattr(paramiko, "SSHClient", lambda: client)
    with pytest.raises(SSHError) as error:
        SSHConnection(Settings("frame", "private password text")).connect()
    assert "private password text" not in str(error.value)


def test_disconnect_releases_even_an_inactive_transport_socket():
    import socket
    left, right = socket.socketpair()
    try:
        connection = SSHConnection(Settings('frame', 'private'))
        transport = paramiko.Transport(left)
        connection.client._transport = transport
        assert not transport.is_active()
        connection.close()
        assert left.fileno() == -1
        assert right.recv(1) == b''
        assert connection.client.get_transport() is None
        connection.close()  # Closing an already completed setup is safe.
    finally:
        left.close()
        right.close()


def test_disconnect_still_closes_transport_when_a_channel_close_fails(monkeypatch):
    client = Mock()
    transport = client.get_transport.return_value
    monkeypatch.setattr(paramiko, 'SSHClient', lambda: client)
    connection = SSHConnection(Settings('frame', 'private'))
    command = Mock()
    command.close.side_effect = OSError('Already closed')
    sftp = Mock()
    connection._command_channels.add(command)
    connection._sftp_clients.add(sftp)
    connection.close()
    command.close.assert_called_once()
    sftp.close.assert_called_once()
    transport.close.assert_called_once()
    transport.sock.close.assert_called_once()
    client.close.assert_called_once()
    assert not connection._command_channels and not connection._sftp_clients
    connection.close()
    client.close.assert_called_once()


@pytest.mark.parametrize('output,expected', [
    (b'HFI_PROGRESS {"stage":"build","message":"Working"}\nHFI_ERROR Native build failed.\nE: setgroups failed\nprivate\n',
     'Native build failed.\nE: setgroups failed\n[redacted]'),
    (b'HFI_PROGRESS {}\nHFI_RESULT {}\n', 'Remote step failed (exit 1).'),
    (b'HFI_LOG {"stage":"compile","message":"[22/100] task"}\nHFI_ERROR Native build failed.\ncompiler diagnostic private\n',
     'Native build failed.\ncompiler diagnostic [redacted]'),
    (b'connection helper failed\n', 'connection helper failed'),
])
def test_remote_failure_shows_diagnostic_without_progress_protocol(monkeypatch, output, expected):
    client = Mock()
    channel = Mock()
    states = iter([True, False, False])
    channel.recv_ready.side_effect = lambda: next(states)
    channel.recv.return_value = output
    channel.exit_status_ready.return_value = True
    channel.recv_exit_status.return_value = 1
    stdout = Mock(channel=channel)
    client.exec_command.return_value = (Mock(), stdout, Mock())
    monkeypatch.setattr(paramiko, 'SSHClient', lambda: client)
    connection = SSHConnection(Settings('frame', 'private'))
    progress = Mock()
    with pytest.raises(SSHError) as error:
        connection.run(['python3', 'helper.py', 'build'], progress=progress)
    assert str(error.value) == expected
    assert all('HFI_PROGRESS' not in line for line in str(error.value).splitlines())
    assert progress.called
    channel.close.assert_called_once()


def test_long_live_activity_is_forwarded_but_response_memory_is_bounded(monkeypatch):
    from collections import deque
    line = b'HFI_LOG {"stage":"compile","message":"' + b'x' * 1024 + b'"}\n'
    output = line * 800 + b'HFI_RESULT {"built":true}\n'
    packets = deque(output[index:index + 65536] for index in range(0, len(output), 65536))
    channel = Mock()
    channel.recv_ready.side_effect = lambda: bool(packets)
    channel.recv.side_effect = lambda count: packets.popleft()
    channel.exit_status_ready.side_effect = lambda: not packets
    channel.recv_exit_status.return_value = 0
    client = Mock()
    client.exec_command.return_value = (Mock(), Mock(channel=channel), Mock())
    monkeypatch.setattr(paramiko, 'SSHClient', lambda: client)
    progress = Mock()
    result = SSHConnection(Settings('frame', 'private')).run(['python3', 'helper.py', 'build'], progress=progress)
    assert progress.call_count == 801
    assert progress.call_args.args[0] == 'HFI_RESULT {"built":true}'
    assert len(result.encode()) <= 256 * 1024
    assert result.endswith('HFI_RESULT {"built":true}\n')
    channel.close.assert_called_once()


def test_utf8_activity_split_between_ssh_packets_is_preserved(monkeypatch):
    from collections import deque
    output = 'HFI_LOG {"stage":"compile","message":"Compiled café"}\nHFI_RESULT {}\n'.encode()
    split = output.index('é'.encode()) + 1
    packets = deque((output[:split], output[split:]))
    channel = Mock()
    channel.recv_ready.side_effect = lambda: bool(packets)
    channel.recv.side_effect = lambda count: packets.popleft()
    channel.exit_status_ready.side_effect = lambda: not packets
    channel.recv_exit_status.return_value = 0
    client = Mock()
    client.exec_command.return_value = (Mock(), Mock(channel=channel), Mock())
    monkeypatch.setattr(paramiko, 'SSHClient', lambda: client)
    progress = Mock()
    SSHConnection(Settings('frame', 'private')).run(['python3', 'helper.py'], progress=progress)
    assert progress.call_args_list[0].args[0] == 'HFI_LOG {"stage":"compile","message":"Compiled café"}'


def test_uninstall_timeout_reports_uncertain_removal_and_closes_channel(monkeypatch):
    from halo_frame_installer import ssh
    channel = Mock()
    client = Mock()
    client.exec_command.return_value = (Mock(), Mock(channel=channel), Mock())
    monkeypatch.setattr(paramiko, 'SSHClient', lambda: client)
    monkeypatch.setattr(ssh.time, 'monotonic', Mock(side_effect=[0, 601]))
    connection = SSHConnection(Settings('frame', 'private'))
    with pytest.raises(SSHError, match='Removal may be incomplete') as error:
        connection.run(['python3', 'helper.py', 'uninstall'], timeout=600,
                       timeout_message='Removal may be incomplete; check the saved backup and activity log.')
    assert 'were kept' not in str(error.value)
    channel.close.assert_called_once()
    assert not connection._command_channels