diff --git a/ceti/whaletag.py b/ceti/whaletag.py index bd00b55..aaae428 100644 --- a/ceti/whaletag.py +++ b/ceti/whaletag.py @@ -16,6 +16,7 @@ from argparse import Namespace import asyncio import ipaddress +import logging import os import re import socket @@ -27,6 +28,7 @@ from ceti.utils import sha256sum +logger = logging.getLogger(__name__) LOCAL_DATA_PATH = os.path.join(os.getcwd(), "data") DEFAULT_USBGADGET_IPNETWORK = "192.168.11.0/24" @@ -63,7 +65,6 @@ def find_ssh_servers(): # get hostnames for all ssh servers def get_hostname_by_addr(addr): try: - # Connect to the remote whale tag ssh = paramiko.SSHClient() ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) ssh.connect( @@ -74,13 +75,13 @@ def get_hostname_by_addr(addr): hostname = stdout.readline().strip() ssh.close() return hostname - except: + except (paramiko.SSHException, OSError) as e: + logger.warning("Failed to get hostname for %s: %s", addr, e) return "" # Verify we can connect to the remote system using ssh with default credentials def can_connect(addr): try: - # test connecting with ssh using default tag password ssh = paramiko.SSHClient() ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) ssh.connect( @@ -88,7 +89,14 @@ def can_connect(addr): username=DEFAULT_USERNAME, password=DEFAULT_PASSWORD) ssh.close() - except BaseException: + except paramiko.AuthenticationException as e: + logger.error("Authentication failed for %s: %s", addr, e) + return False + except paramiko.SSHException as e: + logger.error("SSH error connecting to %s: %s", addr, e) + return False + except OSError as e: + logger.error("Network error connecting to %s: %s", addr, e) return False return True diff --git a/tests/test_whaletag.py b/tests/test_whaletag.py new file mode 100644 index 0000000..be01d96 --- /dev/null +++ b/tests/test_whaletag.py @@ -0,0 +1,101 @@ +import logging +from unittest import mock + +import paramiko +import pytest + +from ceti.whaletag import can_connect, get_hostname_by_addr + + +class TestCanConnect: + """Tests for can_connect() -- issue #39.""" + + @mock.patch("ceti.whaletag.paramiko.SSHClient") + def test_successful_connection(self, mock_ssh_cls): + assert can_connect("192.168.1.10") is True + mock_ssh_cls.return_value.connect.assert_called_once() + mock_ssh_cls.return_value.close.assert_called_once() + + @mock.patch("ceti.whaletag.paramiko.SSHClient") + def test_auth_failure_returns_false(self, mock_ssh_cls, caplog): + mock_ssh_cls.return_value.connect.side_effect = ( + paramiko.AuthenticationException("Bad password") + ) + with caplog.at_level(logging.ERROR): + result = can_connect("192.168.1.10") + + assert result is False + assert "Authentication failed" in caplog.text + assert "192.168.1.10" in caplog.text + + @mock.patch("ceti.whaletag.paramiko.SSHClient") + def test_ssh_error_returns_false(self, mock_ssh_cls, caplog): + mock_ssh_cls.return_value.connect.side_effect = ( + paramiko.SSHException("Protocol error") + ) + with caplog.at_level(logging.ERROR): + result = can_connect("192.168.1.10") + + assert result is False + assert "SSH error" in caplog.text + + @mock.patch("ceti.whaletag.paramiko.SSHClient") + def test_network_error_returns_false(self, mock_ssh_cls, caplog): + mock_ssh_cls.return_value.connect.side_effect = ( + OSError("Connection refused") + ) + with caplog.at_level(logging.ERROR): + result = can_connect("192.168.1.10") + + assert result is False + assert "Network error" in caplog.text + + @mock.patch("ceti.whaletag.paramiko.SSHClient") + def test_keyboard_interrupt_not_caught(self, mock_ssh_cls): + mock_ssh_cls.return_value.connect.side_effect = KeyboardInterrupt + with pytest.raises(KeyboardInterrupt): + can_connect("192.168.1.10") + + +class TestGetHostnameByAddr: + """Tests for get_hostname_by_addr() -- issue #39.""" + + @mock.patch("ceti.whaletag.paramiko.SSHClient") + def test_returns_hostname(self, mock_ssh_cls): + mock_stdout = mock.Mock() + mock_stdout.readline.return_value = "wt-b827eb123456\n" + mock_ssh_cls.return_value.exec_command.return_value = ( + None, mock_stdout, None + ) + + result = get_hostname_by_addr("192.168.1.10") + assert result == "wt-b827eb123456" + + @mock.patch("ceti.whaletag.paramiko.SSHClient") + def test_connection_failure_returns_empty(self, mock_ssh_cls, caplog): + mock_ssh_cls.return_value.connect.side_effect = ( + paramiko.SSHException("Connection refused") + ) + with caplog.at_level(logging.WARNING): + result = get_hostname_by_addr("192.168.1.10") + + assert result == "" + assert "Failed to get hostname" in caplog.text + assert "192.168.1.10" in caplog.text + + @mock.patch("ceti.whaletag.paramiko.SSHClient") + def test_timeout_returns_empty(self, mock_ssh_cls, caplog): + mock_ssh_cls.return_value.connect.side_effect = ( + OSError("Connection timed out") + ) + with caplog.at_level(logging.WARNING): + result = get_hostname_by_addr("192.168.1.10") + + assert result == "" + assert "Connection timed out" in caplog.text + + @mock.patch("ceti.whaletag.paramiko.SSHClient") + def test_keyboard_interrupt_not_caught(self, mock_ssh_cls): + mock_ssh_cls.return_value.connect.side_effect = KeyboardInterrupt + with pytest.raises(KeyboardInterrupt): + get_hostname_by_addr("192.168.1.10")