rusackas commented on code in PR #40673:
URL: https://github.com/apache/superset/pull/40673#discussion_r3383917239


##########
tests/unit_tests/extensions/ssh_test.py:
##########
@@ -34,3 +64,131 @@ def test_ssh_tunnel_timeout_setting() -> None:
     factory.init_app(app)
     assert sshtunnel.TUNNEL_TIMEOUT == 123.0
     assert sshtunnel.SSH_TIMEOUT == 321.0
+
+
+@patch("superset.extensions.ssh.socket.create_connection")
+@patch("superset.extensions.ssh.paramiko.Transport")
+def test_verify_host_key_match(
+    mock_transport_cls: Mock, mock_create_connection: Mock
+) -> None:
+    # The server presents the same key we expect: verification passes.
+    server_key = paramiko.RSAKey.generate(2048)
+    manager = _make_manager(strict=False)
+    tunnel = _ssh_tunnel(_authorized_key(server_key))
+
+    transport = mock_transport_cls.return_value
+    transport.get_remote_server_key.return_value = server_key
+
+    manager._verify_host_key(tunnel)  # should not raise
+
+    # The TCP connect is bounded by an explicit timeout, and the resulting
+    # socket is handed to Transport.
+    mock_create_connection.assert_called_once_with(
+        ("ssh.example.com", 22), timeout=321.0
+    )
+    
mock_transport_cls.assert_called_once_with(mock_create_connection.return_value)
+    transport.start_client.assert_called_once()
+    transport.close.assert_called_once()
+
+
+@patch("superset.extensions.ssh.socket.create_connection")
+@patch("superset.extensions.ssh.paramiko.Transport")
+def test_verify_host_key_mismatch_raises(
+    mock_transport_cls: Mock, mock_create_connection: Mock
+) -> None:
+    # The server presents a different key than expected: verification fails.
+    expected_key = paramiko.RSAKey.generate(2048)
+    presented_key = paramiko.RSAKey.generate(2048)
+    manager = _make_manager(strict=False)
+    tunnel = _ssh_tunnel(_authorized_key(expected_key))
+
+    transport = mock_transport_cls.return_value
+    transport.get_remote_server_key.return_value = presented_key
+
+    with pytest.raises(SSHTunnelHostKeyVerificationError):
+        manager._verify_host_key(tunnel)
+
+    mock_create_connection.assert_called_once()
+    transport.close.assert_called_once()
+
+
+@patch("superset.extensions.ssh.socket.create_connection")
+def test_verify_host_key_connect_failure_raises(
+    mock_create_connection: Mock,
+) -> None:
+    # A bounded TCP connect failure surfaces as a host-key verification error.
+    manager = _make_manager(strict=False)
+    server_key = paramiko.RSAKey.generate(2048)
+    tunnel = _ssh_tunnel(_authorized_key(server_key))
+
+    mock_create_connection.side_effect = OSError("connection refused")
+
+    with pytest.raises(SSHTunnelHostKeyVerificationError):
+        manager._verify_host_key(tunnel)
+
+
+@patch("superset.extensions.ssh.paramiko.Transport")
+def test_verify_host_key_unset_non_strict_skips(mock_transport_cls: Mock) -> 
None:
+    # Back-compat: no expected key + strict checking off => no verification at 
all.
+    manager = _make_manager(strict=False)
+    tunnel = _ssh_tunnel(None)
+
+    manager._verify_host_key(tunnel)  # should not raise
+
+    mock_transport_cls.assert_not_called()
+
+
+@patch("superset.extensions.ssh.paramiko.Transport")
+def test_verify_host_key_unset_strict_raises(mock_transport_cls: Mock) -> None:
+    # Fail-closed: no expected key + strict checking on => reject.
+    manager = _make_manager(strict=True)
+    tunnel = _ssh_tunnel(None)
+
+    with pytest.raises(SSHTunnelHostKeyVerificationError):
+        manager._verify_host_key(tunnel)
+
+    mock_transport_cls.assert_not_called()
+
+
+@patch("superset.extensions.ssh.socket.create_connection")
+@patch("superset.extensions.ssh.paramiko.Transport")
+def test_verify_host_key_match_ignores_comment_and_whitespace(
+    mock_transport_cls: Mock,
+    mock_create_connection: Mock,
+) -> None:
+    # The stored key may carry a trailing comment and extra whitespace.
+    server_key = paramiko.RSAKey.generate(2048)
+    manager = _make_manager(strict=False)
+    stored = f"  {_authorized_key(server_key)} user@host  "
+    tunnel = _ssh_tunnel(stored)
+
+    transport = mock_transport_cls.return_value
+    transport.get_remote_server_key.return_value = server_key
+
+    manager._verify_host_key(tunnel)  # should not raise
+

Review Comment:
   Good call — added the bounded-connect and Transport handshake assertions 
(socket.create_connection with the configured timeout, Transport built from 
that socket, start_client and close each called once) to 
test_verify_host_key_match_ignores_comment_and_whitespace, matching 
test_verify_host_key_match. A regression that silently skips verification after 
the whitespace/comment stripping would now fail rather than falsely pass. 
Pushed in the latest commit.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to