From 51cf2cdfda5fd4fc67352db329e0d02e05106c8a Mon Sep 17 00:00:00 2001 From: Evan Date: Fri, 5 Jun 2026 10:25:30 -0700 Subject: [PATCH] fix(ssh-tunnel): address re-review feedback on _load_private_key - add circular-import justification comment on the deferred get_default_port import - drop no-op continue in the per-type loop - add NOTE clarifying last_exc reflects the final (RSAKey) attempt - clarify docstring re paramiko 3.2+ PKey.from_path() and why it's unsafe here - strengthen ed25519/RSA tests with isinstance assertions on the parsed key type - add test for passphrase-protected key without password (PasswordRequiredException) - add test for an unparseable key (SSHException listing all attempted types) - de-duplicate test setup via _make_manager / _make_ssh_tunnel helpers - consistent assertion messages across key-type tests Co-Authored-By: Claude Opus 4.8 --- superset/extensions/ssh.py | 16 ++- tests/unit_tests/extensions/ssh_test.py | 132 +++++++++++++++--------- 2 files changed, 96 insertions(+), 52 deletions(-) diff --git a/superset/extensions/ssh.py b/superset/extensions/ssh.py index bebb847ef58..a2076744ed0 100644 --- a/superset/extensions/ssh.py +++ b/superset/extensions/ssh.py @@ -47,9 +47,11 @@ _SSH_KEY_TYPES: tuple[type[PKey], ...] = (Ed25519Key, ECDSAKey, RSAKey) def _load_private_key(pem: str, password: str | None) -> PKey: """Load a private key PEM regardless of algorithm (ed25519, ECDSA, RSA). - paramiko does not expose a polymorphic ``PKey.from_private_key``; each - key class only accepts its own format. Iterate over the supported types - and return the first that parses cleanly. + paramiko 3.2+ has ``PKey.from_path()`` for polymorphic loading, but it + requires a filesystem path; writing private key material to disk would be a + security regression. Each per-class loader only accepts its own format, so + iterate over the supported types on the in-memory ``StringIO`` and return + the first that parses cleanly. """ last_exc: SSHException | None = None for key_class in _SSH_KEY_TYPES: @@ -59,7 +61,10 @@ def _load_private_key(pem: str, password: str | None) -> PKey: raise except SSHException as exc: last_exc = exc - continue + # NOTE: last_exc holds the error from the final attempt (RSAKey), not the + # closest-matching type. For a corrupted ed25519 key, the appended message + # reflects RSAKey's parse error; the full type list above still identifies + # all types attempted. raise SSHException( "Unable to parse SSH private key as any of " f"{', '.join(k.__name__ for k in _SSH_KEY_TYPES)}: {last_exc}" @@ -88,6 +93,9 @@ class SSHManager: ssh_tunnel: "SSHTunnel", sqlalchemy_database_uri: str, ) -> sshtunnel.SSHTunnelForwarder: + # Deferred import to break a circular import: + # superset.utils.ssh_tunnel -> superset.databases.ssh_tunnel.models + # -> superset.extensions -> superset.extensions.ssh (this module). from superset.utils.ssh_tunnel import get_default_port url = make_url_safe(sqlalchemy_database_uri) diff --git a/tests/unit_tests/extensions/ssh_test.py b/tests/unit_tests/extensions/ssh_test.py index 678055430cc..5852cb4c8fc 100644 --- a/tests/unit_tests/extensions/ssh_test.py +++ b/tests/unit_tests/extensions/ssh_test.py @@ -16,6 +16,7 @@ # under the License. from unittest.mock import Mock, patch +import pytest import sshtunnel from cryptography.hazmat.primitives.asymmetric import ec from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey @@ -23,14 +24,37 @@ from cryptography.hazmat.primitives.asymmetric.rsa import ( generate_private_key as generate_rsa_key, ) from cryptography.hazmat.primitives.serialization import ( + BestAvailableEncryption, Encoding, NoEncryption, PrivateFormat, ) +from paramiko import Ed25519Key, PasswordRequiredException, RSAKey, SSHException from superset.extensions.ssh import SSHManager, SSHManagerFactory +def _make_manager() -> SSHManager: + app = Mock() + app.config = { + "SSH_TUNNEL_LOCAL_BIND_ADDRESS": "127.0.0.1", + "SSH_TUNNEL_TIMEOUT_SEC": 10.0, + "SSH_TUNNEL_PACKET_TIMEOUT_SEC": 10.0, + } + return SSHManager(app) + + +def _make_ssh_tunnel(private_key: str, private_key_password: str | None = None) -> Mock: + ssh_tunnel = Mock() + ssh_tunnel.server_address = "ssh.example.com" + ssh_tunnel.server_port = 22 + ssh_tunnel.username = "tunneluser" + ssh_tunnel.password = None + ssh_tunnel.private_key = private_key + ssh_tunnel.private_key_password = private_key_password + return ssh_tunnel + + def test_ssh_tunnel_timeout_setting() -> None: app = Mock() app.config = { @@ -72,21 +96,8 @@ def test_create_tunnel_accepts_ed25519_private_key() -> None: test does not actually open a network connection — only the key parsing path is exercised. """ - app = Mock() - app.config = { - "SSH_TUNNEL_LOCAL_BIND_ADDRESS": "127.0.0.1", - "SSH_TUNNEL_TIMEOUT_SEC": 10.0, - "SSH_TUNNEL_PACKET_TIMEOUT_SEC": 10.0, - } - manager = SSHManager(app) - - ssh_tunnel = Mock() - ssh_tunnel.server_address = "ssh.example.com" - ssh_tunnel.server_port = 22 - ssh_tunnel.username = "tunneluser" - ssh_tunnel.password = None - ssh_tunnel.private_key = _make_ed25519_pem() - ssh_tunnel.private_key_password = None + manager = _make_manager() + ssh_tunnel = _make_ssh_tunnel(_make_ed25519_pem()) with patch("superset.extensions.ssh.sshtunnel.open_tunnel") as mock_open: manager.create_tunnel( @@ -96,7 +107,7 @@ def test_create_tunnel_accepts_ed25519_private_key() -> None: # Key-type-agnostic loader must produce a paramiko PKey usable as ssh_pkey. assert mock_open.called, "open_tunnel was never invoked — key parsing aborted" forwarded_pkey = mock_open.call_args.kwargs["ssh_pkey"] - assert forwarded_pkey is not None + assert isinstance(forwarded_pkey, Ed25519Key) def test_create_tunnel_accepts_rsa_private_key_unchanged() -> None: @@ -116,29 +127,16 @@ def test_create_tunnel_accepts_rsa_private_key_unchanged() -> None: .decode() ) - app = Mock() - app.config = { - "SSH_TUNNEL_LOCAL_BIND_ADDRESS": "127.0.0.1", - "SSH_TUNNEL_TIMEOUT_SEC": 10.0, - "SSH_TUNNEL_PACKET_TIMEOUT_SEC": 10.0, - } - manager = SSHManager(app) - - ssh_tunnel = Mock() - ssh_tunnel.server_address = "ssh.example.com" - ssh_tunnel.server_port = 22 - ssh_tunnel.username = "tunneluser" - ssh_tunnel.password = None - ssh_tunnel.private_key = rsa_pem - ssh_tunnel.private_key_password = None + manager = _make_manager() + ssh_tunnel = _make_ssh_tunnel(rsa_pem) with patch("superset.extensions.ssh.sshtunnel.open_tunnel") as mock_open: manager.create_tunnel( ssh_tunnel, "postgresql://user:pass@db.example.com:5432/x" ) - assert mock_open.called - assert mock_open.call_args.kwargs["ssh_pkey"] is not None + assert mock_open.called, "open_tunnel was never invoked — RSA key parsing aborted" + assert isinstance(mock_open.call_args.kwargs["ssh_pkey"], RSAKey) def test_create_tunnel_accepts_ecdsa_private_key() -> None: @@ -157,21 +155,8 @@ def test_create_tunnel_accepts_ecdsa_private_key() -> None: .decode() ) - app = Mock() - app.config = { - "SSH_TUNNEL_LOCAL_BIND_ADDRESS": "127.0.0.1", - "SSH_TUNNEL_TIMEOUT_SEC": 10.0, - "SSH_TUNNEL_PACKET_TIMEOUT_SEC": 10.0, - } - manager = SSHManager(app) - - ssh_tunnel = Mock() - ssh_tunnel.server_address = "ssh.example.com" - ssh_tunnel.server_port = 22 - ssh_tunnel.username = "tunneluser" - ssh_tunnel.password = None - ssh_tunnel.private_key = ecdsa_pem - ssh_tunnel.private_key_password = None + manager = _make_manager() + ssh_tunnel = _make_ssh_tunnel(ecdsa_pem) with patch("superset.extensions.ssh.sshtunnel.open_tunnel") as mock_open: manager.create_tunnel( @@ -180,3 +165,54 @@ def test_create_tunnel_accepts_ecdsa_private_key() -> None: assert mock_open.called, "open_tunnel was never invoked — ECDSA key parsing aborted" assert mock_open.call_args.kwargs["ssh_pkey"] is not None + + +def test_create_tunnel_passphrase_protected_key_without_password() -> None: + """ + A passphrase-protected key supplied without a passphrase must surface as + ``PasswordRequiredException`` (an actionable "key requires passphrase" + signal) rather than being absorbed by the per-type loop and reported as a + generic "Unable to parse" error. + """ + encrypted_pem = ( + Ed25519PrivateKey.generate() + .private_bytes( + encoding=Encoding.PEM, + format=PrivateFormat.OpenSSH, + encryption_algorithm=BestAvailableEncryption(b"correct horse"), + ) + .decode() + ) + + manager = _make_manager() + ssh_tunnel = _make_ssh_tunnel(encrypted_pem, private_key_password=None) + + with patch("superset.extensions.ssh.sshtunnel.open_tunnel") as mock_open: + with pytest.raises(PasswordRequiredException): + manager.create_tunnel( + ssh_tunnel, "postgresql://user:pass@db.example.com:5432/x" + ) + + assert not mock_open.called + + +def test_create_tunnel_invalid_key_raises_combined_error() -> None: + """ + When a key parses as none of the supported types, ``_load_private_key`` + raises ``SSHException`` whose message lists every type that was attempted, + so the failure clearly communicates that all loaders were tried. + """ + manager = _make_manager() + ssh_tunnel = _make_ssh_tunnel("not a valid private key") + + with patch("superset.extensions.ssh.sshtunnel.open_tunnel") as mock_open: + with pytest.raises(SSHException) as exc_info: + manager.create_tunnel( + ssh_tunnel, "postgresql://user:pass@db.example.com:5432/x" + ) + + message = str(exc_info.value) + assert "Ed25519Key" in message + assert "ECDSAKey" in message + assert "RSAKey" in message + assert not mock_open.called