Add exception to catch single line private keys (#23043)

apache · May 9, 2022 · e63dbdc · e63dbdc
1 parent d7b85d9
commit e63dbdc
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/airflow/providers/ssh/hooks/ssh.py b/airflow/providers/ssh/hooks/ssh.py
@@ -415,6 +415,9 @@ def _pkey_from_private_key(self, private_key: str, passphrase: Optional[str] = N
         :return: ``paramiko.PKey`` appropriate for given key
         :raises AirflowException: if key cannot be read
         """
+        if len(private_key.split("\n", 2)) < 2:
+            raise AirflowException('Key must have BEGIN and END header/footer on separate lines.')
+
         for pkey_class in self._pkey_loaders:
             try:
                 key = pkey_class.from_private_key(StringIO(private_key), password=passphrase)

diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
@@ -740,6 +740,24 @@ def test_openssh_private_key(self):
             session.delete(conn)
             session.commit()
 
+    def test_oneline_key(self):
+        with pytest.raises(Exception):
+            TEST_ONELINE_KEY = "-----BEGIN OPENSSH" + "PRIVATE KEY-----asdfg-----END OPENSSH PRIVATE KEY-----"
+            session = settings.Session()
+            try:
+                conn = Connection(
+                    conn_id='openssh_pkey',
+                    host='localhost',
+                    conn_type='ssh',
+                    extra={"private_key": TEST_ONELINE_KEY},
+                )
+                session.add(conn)
+                session.flush()
+                SSHHook(ssh_conn_id=conn.conn_id)
+            finally:
+                session.delete(conn)
+                session.commit()
+
     @pytest.mark.flaky(max_runs=5, min_passes=1)
     def test_exec_ssh_client_command(self):
         hook = SSHHook(