From 6f6b7956d1eb37cde91c29ada7216dbcf5ff435f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Thu, 17 Dec 2020 01:04:09 +0100 Subject: [PATCH 1/3] Add system tests for CloudKMSHook --- airflow/providers/google/cloud/hooks/kms.py | 3 +- .../google/cloud/hooks/test_kms_hook.py | 104 ++++++++++++++++++ .../google/cloud/utils/gcp_authenticator.py | 1 + tests/test_utils/gcp_system_helpers.py | 3 + 4 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 tests/providers/google/cloud/hooks/test_kms_hook.py diff --git a/airflow/providers/google/cloud/hooks/kms.py b/airflow/providers/google/cloud/hooks/kms.py index 00f7216b99e07..e63c2f1cb79f6 100644 --- a/airflow/providers/google/cloud/hooks/kms.py +++ b/airflow/providers/google/cloud/hooks/kms.py @@ -169,5 +169,4 @@ def decrypt( metadata=metadata, ) - plaintext = response.plaintext - return plaintext + return response.plaintext diff --git a/tests/providers/google/cloud/hooks/test_kms_hook.py b/tests/providers/google/cloud/hooks/test_kms_hook.py new file mode 100644 index 0000000000000..5c77ed859f928 --- /dev/null +++ b/tests/providers/google/cloud/hooks/test_kms_hook.py @@ -0,0 +1,104 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import base64 +import os +from tempfile import TemporaryDirectory + +import pytest + +from airflow.providers.google.cloud.hooks.kms import CloudKMSHook +from tests.providers.google.cloud.utils.gcp_authenticator import GCP_KMS_KEY +from tests.test_utils.gcp_system_helpers import GoogleSystemTest, provide_gcp_context + +# To prevent resource name collisions, key ring and key resources CANNOT be deleted, so +# to avoid cluttering the project, we only create the key once during project initialization. +# See: https://cloud.google.com/kms/docs/faq#cannot_delete +GCP_KMS_KEYRING_NAME = os.environ.get('GCP_KMS_KEYRING_NAME', 'test-airflow-system-tests-keyring') +GCP_KMS_KEY_NAME = os.environ.get('GCP_KMS_KEY_NAME', 'test-airflow-system-tests-key') + + +@pytest.mark.credential_file(GCP_KMS_KEY) +class TestKmsHook(GoogleSystemTest): + @provide_gcp_context(GCP_KMS_KEY) + def test_encrypt(self): + with TemporaryDirectory() as tmp_dir: + kms_hook = CloudKMSHook() + content = kms_hook.encrypt( + key_name=( + f"projects/{kms_hook.project_id}/locations/global/keyRings/" + f"{GCP_KMS_KEYRING_NAME}/cryptoKeys/{GCP_KMS_KEY_NAME}" + ), + plaintext=b"TEST-SECRET", + ) + with open(f"{tmp_dir}/mysecret.txt.encrypted", "wb") as encrypted_file: + encrypted_file.write(base64.b64decode(content)) + self.execute_cmd( + [ + "gcloud", + "kms", + "decrypt", + "--location", + "global", + "--keyring", + GCP_KMS_KEYRING_NAME, + "--key", + GCP_KMS_KEY_NAME, + "--ciphertext-file", + f"{tmp_dir}/mysecret.txt.encrypted", + "--plaintext-file", + f"{tmp_dir}/mysecret.txt", + ] + ) + with open(f"{tmp_dir}/mysecret.txt", "rb") as secret_file: + secret = secret_file.read() + self.assertEqual(secret, b"TEST-SECRET") + + @provide_gcp_context(GCP_KMS_KEY) + def test_decrypt(self): + with TemporaryDirectory() as tmp_dir: + with open(f"{tmp_dir}/mysecret.txt", "w") as secret_file: + secret_file.write("TEST-SECRET") + self.execute_cmd( + [ + "gcloud", + "kms", + "encrypt", + "--location", + "global", + "--keyring", + GCP_KMS_KEYRING_NAME, + "--key", + GCP_KMS_KEY_NAME, + "--plaintext-file", + f"{tmp_dir}/mysecret.txt", + "--ciphertext-file", + f"{tmp_dir}/mysecret.txt.encrypted", + ] + ) + with open(f"{tmp_dir}/mysecret.txt.encrypted", "rb") as encrypted_file: + encrypted_secret = base64.b64encode(encrypted_file.read()).decode() + + kms_hook = CloudKMSHook() + content = kms_hook.decrypt( + key_name=( + f"projects/{kms_hook.project_id}/locations/global/keyRings/" + f"{GCP_KMS_KEYRING_NAME}/cryptoKeys/{GCP_KMS_KEY_NAME}" + ), + ciphertext=encrypted_secret, + ) + self.assertEqual(content, b"TEST-SECRET") diff --git a/tests/providers/google/cloud/utils/gcp_authenticator.py b/tests/providers/google/cloud/utils/gcp_authenticator.py index 945dd3353bf4f..bf36eadc2202b 100644 --- a/tests/providers/google/cloud/utils/gcp_authenticator.py +++ b/tests/providers/google/cloud/utils/gcp_authenticator.py @@ -46,6 +46,7 @@ GCP_GCS_KEY = 'gcp_gcs.json' GCP_GCS_TRANSFER_KEY = 'gcp_gcs_transfer.json' GCP_GKE_KEY = "gcp_gke.json" +GCP_KMS_KEY = "gcp_kms.json" GCP_LIFE_SCIENCES_KEY = 'gcp_life_sciences.json' GCP_MEMORYSTORE = 'gcp_memorystore.json' GCP_PUBSUB_KEY = "gcp_pubsub.json" diff --git a/tests/test_utils/gcp_system_helpers.py b/tests/test_utils/gcp_system_helpers.py index 250593af240d0..65721111725c5 100644 --- a/tests/test_utils/gcp_system_helpers.py +++ b/tests/test_utils/gcp_system_helpers.py @@ -83,9 +83,12 @@ def provide_gcp_context( :param scopes: OAuth scopes for the connection :type scopes: Sequence :param project_id: The id of GCP project for the connection. + Default: ``os.environ["GCP_PROJECT_ID"]`` or None :type project_id: str """ key_file_path = resolve_full_gcp_key_path(key_file_path) # type: ignore + if project_id is None: + project_id = os.environ.get("GCP_PROJECT_ID") with provide_gcp_conn_and_credentials( key_file_path, scopes, project_id ), tempfile.TemporaryDirectory() as gcloud_config_tmp, mock.patch.dict( From 6e4c057fd03cae59c6fd15ba8bfd9f9b5888c2bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Thu, 17 Dec 2020 11:03:56 +0100 Subject: [PATCH 2/3] fixup! Add system tests for CloudKMSHook --- .../google/cloud/hooks/{test_kms_hook.py => test_kms_system.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/providers/google/cloud/hooks/{test_kms_hook.py => test_kms_system.py} (100%) diff --git a/tests/providers/google/cloud/hooks/test_kms_hook.py b/tests/providers/google/cloud/hooks/test_kms_system.py similarity index 100% rename from tests/providers/google/cloud/hooks/test_kms_hook.py rename to tests/providers/google/cloud/hooks/test_kms_system.py From 1a83fe14377db2a96cc9f1d440e84eb7827e222b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Thu, 17 Dec 2020 01:25:29 +0100 Subject: [PATCH 3/3] Update compatibility with google-cloud-kms>=2.0 --- airflow/providers/google/cloud/hooks/kms.py | 20 ++++++---- setup.py | 2 +- .../providers/google/cloud/hooks/test_kms.py | 40 +++++++++++-------- 3 files changed, 37 insertions(+), 25 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/kms.py b/airflow/providers/google/cloud/hooks/kms.py index e63c2f1cb79f6..3fd14336d23d7 100644 --- a/airflow/providers/google/cloud/hooks/kms.py +++ b/airflow/providers/google/cloud/hooks/kms.py @@ -118,12 +118,14 @@ def encrypt( :rtype: str """ response = self.get_conn().encrypt( - name=key_name, - plaintext=plaintext, - additional_authenticated_data=authenticated_data, + request={ + 'name': key_name, + 'plaintext': plaintext, + 'additional_authenticated_data': authenticated_data, + }, retry=retry, timeout=timeout, - metadata=metadata, + metadata=metadata or (), ) ciphertext = _b64encode(response.ciphertext) @@ -161,12 +163,14 @@ def decrypt( :rtype: bytes """ response = self.get_conn().decrypt( - name=key_name, - ciphertext=_b64decode(ciphertext), - additional_authenticated_data=authenticated_data, + request={ + 'name': key_name, + 'ciphertext': _b64decode(ciphertext), + 'additional_authenticated_data': authenticated_data, + }, retry=retry, timeout=timeout, - metadata=metadata, + metadata=metadata or (), ) return response.plaintext diff --git a/setup.py b/setup.py index 45ebfa1846e41..501fee6329989 100644 --- a/setup.py +++ b/setup.py @@ -260,7 +260,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version 'google-cloud-datacatalog>=0.5.0, <0.8', # TODO: we should migrate to 1.0 likely and add <2.0.0 then 'google-cloud-dataproc>=1.0.1,<2.0.0', 'google-cloud-dlp>=0.11.0,<2.0.0', - 'google-cloud-kms>=1.2.1,<2.0.0', + 'google-cloud-kms>=2.0.0,<3.0.0', 'google-cloud-language>=1.1.1,<2.0.0', 'google-cloud-logging>=1.14.0,<2.0.0', 'google-cloud-memcache>=0.2.0', diff --git a/tests/providers/google/cloud/hooks/test_kms.py b/tests/providers/google/cloud/hooks/test_kms.py index 894532aaed41b..d2f4519df5534 100644 --- a/tests/providers/google/cloud/hooks/test_kms.py +++ b/tests/providers/google/cloud/hooks/test_kms.py @@ -82,12 +82,14 @@ def test_encrypt(self, mock_get_conn): result = self.kms_hook.encrypt(TEST_KEY_ID, PLAINTEXT) mock_get_conn.assert_called_once_with() mock_get_conn.return_value.encrypt.assert_called_once_with( - name=TEST_KEY_ID, - plaintext=PLAINTEXT, - additional_authenticated_data=None, + request=dict( + name=TEST_KEY_ID, + plaintext=PLAINTEXT, + additional_authenticated_data=None, + ), retry=None, timeout=None, - metadata=None, + metadata=(), ) self.assertEqual(PLAINTEXT_b64, result) @@ -97,12 +99,14 @@ def test_encrypt_with_auth_data(self, mock_get_conn): result = self.kms_hook.encrypt(TEST_KEY_ID, PLAINTEXT, AUTH_DATA) mock_get_conn.assert_called_once_with() mock_get_conn.return_value.encrypt.assert_called_once_with( - name=TEST_KEY_ID, - plaintext=PLAINTEXT, - additional_authenticated_data=AUTH_DATA, + request=dict( + name=TEST_KEY_ID, + plaintext=PLAINTEXT, + additional_authenticated_data=AUTH_DATA, + ), retry=None, timeout=None, - metadata=None, + metadata=(), ) self.assertEqual(PLAINTEXT_b64, result) @@ -112,12 +116,14 @@ def test_decrypt(self, mock_get_conn): result = self.kms_hook.decrypt(TEST_KEY_ID, CIPHERTEXT_b64) mock_get_conn.assert_called_once_with() mock_get_conn.return_value.decrypt.assert_called_once_with( - name=TEST_KEY_ID, - ciphertext=CIPHERTEXT, - additional_authenticated_data=None, + request=dict( + name=TEST_KEY_ID, + ciphertext=CIPHERTEXT, + additional_authenticated_data=None, + ), retry=None, timeout=None, - metadata=None, + metadata=(), ) self.assertEqual(PLAINTEXT, result) @@ -127,11 +133,13 @@ def test_decrypt_with_auth_data(self, mock_get_conn): result = self.kms_hook.decrypt(TEST_KEY_ID, CIPHERTEXT_b64, AUTH_DATA) mock_get_conn.assert_called_once_with() mock_get_conn.return_value.decrypt.assert_called_once_with( - name=TEST_KEY_ID, - ciphertext=CIPHERTEXT, - additional_authenticated_data=AUTH_DATA, + request=dict( + name=TEST_KEY_ID, + ciphertext=CIPHERTEXT, + additional_authenticated_data=AUTH_DATA, + ), retry=None, timeout=None, - metadata=None, + metadata=(), ) self.assertEqual(PLAINTEXT, result)