Added support for DSS, ECDSA, and Ed25519 private keys in SSHHook (#12467)

Author: brighton1101

airflow/providers/ssh/hooks/ssh.py

@@ -58,6 +58,14 @@ class SSHHook(BaseHook):  # pylint: disable=too-many-instance-attributes
     :type keepalive_interval: int
     """
 
+    # key type name to paramiko PKey class
+    _default_pkey_mappings = {
+        'dsa': paramiko.DSSKey,
+        'ecdsa': paramiko.ECDSAKey,
+        'ed25519': paramiko.Ed25519Key,
+        'rsa': paramiko.RSAKey,
+    }
+
     conn_name_attr = 'ssh_conn_id'
     default_conn_name = 'ssh_default'
     conn_type = 'ssh'
@@ -124,13 +132,8 @@ def __init__(  # pylint: disable=too-many-statements
 
                 private_key = extra_options.get('private_key')
                 private_key_passphrase = extra_options.get('private_key_passphrase')
-                if private_key and private_key_passphrase:
-                    self.pkey = paramiko.RSAKey.from_private_key(
-                        StringIO(private_key), password=private_key_passphrase
-                    )
-                elif private_key and not private_key_passphrase:
-                    self.pkey = paramiko.RSAKey.from_private_key(StringIO(private_key))
-
+                if private_key:
+                    self.pkey = self._pkey_from_private_key(private_key, passphrase=private_key_passphrase)
                 if "timeout" in extra_options:
                     self.timeout = int(extra_options["timeout"], 10)
 
@@ -320,3 +323,24 @@ def create_tunnel(
         )
 
         return self.get_tunnel(remote_port, remote_host, local_port)
+
+    def _pkey_from_private_key(self, private_key: str, passphrase: Optional[str] = None) -> paramiko.PKey:
+        """
+        Creates appropriate paramiko key for given private key
+
+        :param private_key: string containing private key
+        :return: `paramiko.PKey` appropriate for given key
+        :raises AirflowException: if key cannot be read
+        """
+        allowed_pkey_types = self._default_pkey_mappings.values()
+        for pkey_type in allowed_pkey_types:
+            try:
+                key = pkey_type.from_private_key(StringIO(private_key), password=passphrase)
+                return key
+            except paramiko.ssh_exception.SSHException:
+                continue
+        raise AirflowException(
+            'Private key provided cannot be read by paramiko.'
+            'Ensure key provided is valid for one of the following'
+            'key formats: RSA, DSS, ECDSA, or Ed25519'
+        )

tests/providers/ssh/hooks/test_ssh.py

@@ -63,12 +63,16 @@ def generate_host_key(pkey: paramiko.PKey):
 TEST_PRIVATE_KEY = generate_key_string(pkey=TEST_PKEY)
 TEST_HOST_KEY = generate_host_key(pkey=TEST_PKEY)
 
+TEST_PKEY_ECDSA = paramiko.ECDSAKey.generate()
+TEST_PRIVATE_KEY_ECDSA = generate_key_string(pkey=TEST_PKEY_ECDSA)
+
 PASSPHRASE = ''.join(random.choice(string.ascii_letters) for i in range(10))
 TEST_ENCRYPTED_PRIVATE_KEY = generate_key_string(pkey=TEST_PKEY, passphrase=PASSPHRASE)
 
 
 class TestSSHHook(unittest.TestCase):
     CONN_SSH_WITH_PRIVATE_KEY_EXTRA = 'ssh_with_private_key_extra'
+    CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA = 'ssh_with_private_key_ecdsa_extra'
     CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA = 'ssh_with_private_key_passphrase_extra'
     CONN_SSH_WITH_EXTRA = 'ssh_with_extra'
     CONN_SSH_WITH_EXTRA_FALSE_LOOK_FOR_KEYS = 'ssh_with_extra_false_look_for_keys'
@@ -83,6 +87,7 @@ def tearDownClass(cls) -> None:
             conns_to_reset = [
                 cls.CONN_SSH_WITH_PRIVATE_KEY_EXTRA,
                 cls.CONN_SSH_WITH_PRIVATE_KEY_PASSPHRASE_EXTRA,
+                cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
                 cls.CONN_SSH_WITH_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
                 cls.CONN_SSH_WITH_HOST_KEY_AND_NO_HOST_KEY_CHECK_FALSE,
@@ -134,6 +139,14 @@ def setUpClass(cls) -> None:
                 ),
             )
         )
+        db.merge_conn(
+            Connection(
+                conn_id=cls.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
+                host='localhost',
+                conn_type='ssh',
+                extra=json.dumps({"private_key": TEST_PRIVATE_KEY_ECDSA}),
+            )
+        )
         db.merge_conn(
             Connection(
                 conn_id=cls.CONN_SSH_WITH_HOST_KEY_EXTRA,
@@ -313,6 +326,29 @@ def test_tunnel_with_private_key_passphrase(self, ssh_mock):
                 logger=hook.log,
             )
 
+    @mock.patch('airflow.providers.ssh.hooks.ssh.SSHTunnelForwarder')
+    def test_tunnel_with_private_key_ecdsa(self, ssh_mock):
+        hook = SSHHook(
+            ssh_conn_id=self.CONN_SSH_WITH_PRIVATE_KEY_ECDSA_EXTRA,
+            remote_host='remote_host',
+            port='port',
+            username='username',
+            timeout=10,
+        )
+
+        with hook.get_tunnel(1234):
+            ssh_mock.assert_called_once_with(
+                'remote_host',
+                ssh_port='port',
+                ssh_username='username',
+                ssh_pkey=TEST_PKEY_ECDSA,
+                ssh_proxy=None,
+                local_bind_address=('localhost',),
+                remote_bind_address=('localhost', 1234),
+                host_pkey_directories=[],
+                logger=hook.log,
+            )
+
     def test_ssh_connection(self):
         hook = SSHHook(ssh_conn_id='ssh_default')
         with hook.get_conn() as client: