From a6161ab3456b28a86b371056751ee2f05942c885 Mon Sep 17 00:00:00 2001 From: Francis Date: Fri, 17 Apr 2026 16:30:52 -0400 Subject: [PATCH] Replace bare exception handlers with specific paramiko/OS exceptions get_hostname_by_addr() and can_connect() were catching all exceptions silently, making it impossible to debug connection failures in the field. Now catches paramiko.SSHException, AuthenticationException, and OSError specifically, and logs the error with the target address. Adds 9 unit tests covering success, auth failure, SSH errors, network errors, and KeyboardInterrupt propagation. Fixes #39 --- ceti/whaletag.py | 16 +++++-- tests/test_whaletag.py | 101 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 4 deletions(-) create mode 100644 tests/test_whaletag.py 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")