diff --git a/providers/teradata/src/airflow/providers/teradata/utils/encryption_utils.py b/providers/teradata/src/airflow/providers/teradata/utils/encryption_utils.py index 57ed4b9855810..3ce50d0070e17 100644 --- a/providers/teradata/src/airflow/providers/teradata/utils/encryption_utils.py +++ b/providers/teradata/src/airflow/providers/teradata/utils/encryption_utils.py @@ -21,6 +21,8 @@ import string import subprocess +from airflow.providers.teradata.utils.tpt_util import get_remote_os + def generate_random_password(length=12): # Define the character set: letters, digits, and special characters @@ -51,25 +53,27 @@ def generate_encrypted_file_with_openssl(file_path: str, password: str, out_file def decrypt_remote_file_to_string(ssh_client, remote_enc_file, password, bteq_command_str): - # Run openssl decrypt command on remote machine - quoted_password = shell_quote_single(password) + remote_os = get_remote_os(ssh_client) + windows_remote = remote_os == "windows" + + if windows_remote: + # Windows cmd.exe does not support single quotes; use double quotes with "" escaping + password_escaped = password.replace('"', '""') + quoted_password = f'"{password_escaped}"' + else: + # Unix shells require closing the single quote, adding an escaped quote, then reopening + password_escaped = password.replace("'", "'\\''") + quoted_password = f"'{password_escaped}'" decrypt_cmd = ( f"openssl enc -d -aes-256-cbc -salt -pbkdf2 -pass pass:{quoted_password} -in {remote_enc_file} | " + bteq_command_str ) - # Clear password to prevent lingering sensitive data + # Clear sensitive data from memory password = None quoted_password = None stdin, stdout, stderr = ssh_client.exec_command(decrypt_cmd) - # Wait for command to finish exit_status = stdout.channel.recv_exit_status() output = stdout.read().decode() err = stderr.read().decode() return exit_status, output, err - - -def shell_quote_single(s): - # Escape single quotes in s, then wrap in single quotes - # In shell, to include a single quote inside single quotes, close, add '\'' and reopen - return "'" + s.replace("'", "'\\''") + "'" diff --git a/providers/teradata/tests/unit/teradata/utils/test_encryption_utils.py b/providers/teradata/tests/unit/teradata/utils/test_encryption_utils.py index fdf1c8bc21b82..8562da9dde1d6 100644 --- a/providers/teradata/tests/unit/teradata/utils/test_encryption_utils.py +++ b/providers/teradata/tests/unit/teradata/utils/test_encryption_utils.py @@ -23,7 +23,6 @@ decrypt_remote_file_to_string, generate_encrypted_file_with_openssl, generate_random_password, - shell_quote_single, ) @@ -31,7 +30,6 @@ class TestEncryptionUtils: def test_generate_random_password_length(self): pwd = generate_random_password(16) assert len(pwd) == 16 - # Check characters are in allowed set allowed_chars = string.ascii_letters + string.digits + string.punctuation assert (all(c in allowed_chars for c in pwd)) is True @@ -60,43 +58,57 @@ def test_generate_encrypted_file_with_openssl_calls_subprocess(self, mock_run): check=True, ) - def test_shell_quote_single_simple(self): - s = "simple" - quoted = shell_quote_single(s) - assert quoted == "'simple'" - - def test_shell_quote_single_with_single_quote(self): - s = "O'Reilly" - quoted = shell_quote_single(s) - assert quoted == "'O'\\''Reilly'" - - def test_decrypt_remote_file_to_string(self): + @patch("airflow.providers.teradata.utils.encryption_utils.get_remote_os") + def test_decrypt_remote_file_to_string_unix(self, mock_get_remote_os): + mock_get_remote_os.return_value = "unix" password = "mysecret" remote_enc_file = "/remote/encrypted.enc" bteq_command_str = "bteq -c UTF-8" ssh_client = MagicMock() - mock_stdin = MagicMock() mock_stdout = MagicMock() + mock_stdout.channel.recv_exit_status.return_value = 0 + mock_stdout.read.return_value = b"decrypted output" mock_stderr = MagicMock() + mock_stderr.read.return_value = b"" + ssh_client.exec_command.return_value = (MagicMock(), mock_stdout, mock_stderr) + + exit_status, output, err = decrypt_remote_file_to_string( + ssh_client, remote_enc_file, password, bteq_command_str + ) + + expected_cmd = ( + f"openssl enc -d -aes-256-cbc -salt -pbkdf2 -pass pass:'mysecret' -in {remote_enc_file} | " + + bteq_command_str + ) + ssh_client.exec_command.assert_called_once_with(expected_cmd) + assert exit_status == 0 + assert output == "decrypted output" + assert err == "" - # Setup mock outputs and exit code + @patch("airflow.providers.teradata.utils.encryption_utils.get_remote_os") + def test_decrypt_remote_file_to_string_windows(self, mock_get_remote_os): + mock_get_remote_os.return_value = "windows" + password = "mysecret" + remote_enc_file = "/remote/encrypted.enc" + bteq_command_str = "bteq -c UTF-8" + + ssh_client = MagicMock() + mock_stdout = MagicMock() mock_stdout.channel.recv_exit_status.return_value = 0 mock_stdout.read.return_value = b"decrypted output" + mock_stderr = MagicMock() mock_stderr.read.return_value = b"" - - ssh_client.exec_command.return_value = (mock_stdin, mock_stdout, mock_stderr) + ssh_client.exec_command.return_value = (MagicMock(), mock_stdout, mock_stderr) exit_status, output, err = decrypt_remote_file_to_string( ssh_client, remote_enc_file, password, bteq_command_str ) - quoted_password = shell_quote_single(password) expected_cmd = ( - f"openssl enc -d -aes-256-cbc -salt -pbkdf2 -pass pass:{quoted_password} -in {remote_enc_file} | " + f'openssl enc -d -aes-256-cbc -salt -pbkdf2 -pass pass:"mysecret" -in {remote_enc_file} | ' + bteq_command_str ) - ssh_client.exec_command.assert_called_once_with(expected_cmd) assert exit_status == 0 assert output == "decrypted output"