Replace bare exception handlers in whaletag SSH functions

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,21 +75,28 @@ 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(
             hostname=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
 

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")