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()
assert secret == b"TEST-SECRET"
class TestCloudKMSHook(unittest.TestCase):
def setUp(self):
with mock.patch(
"airflow.providers.google.cloud.hooks.base.CloudBaseHook.__init__",
new=mock_init,
):
self.kms_hook = CloudKMSHook(gcp_conn_id="test")
@mock.patch(
"airflow.providers.google.cloud.hooks.kms.CloudKMSHook.client_info",
new_callable=mock.PropertyMock,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.kms.CloudKMSHook._get_credentials"
)
@mock.patch(
"airflow.providers.google.cloud.hooks.kms.KeyManagementServiceClient")
def test_kms_client_creation(self, mock_client, mock_get_creds,
mock_client_info):
result = self.kms_hook.get_conn()
mock_client.assert_called_once_with(
credentials=mock_get_creds.return_value,
client_info=mock_client_info.return_value,
)
self.assertEqual(mock_client.return_value, result)
self.assertEqual(self.kms_hook._conn, result)
@mock.patch(
"airflow.providers.google.cloud.hooks.kms.CloudKMSHook.get_conn")
def test_encrypt(self, mock_get_conn):
mock_get_conn.return_value.encrypt.return_value = RESPONSE
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,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual(PLAINTEXT_b64, result)
@mock.patch(
"airflow.providers.google.cloud.hooks.kms.CloudKMSHook.get_conn")
def test_encrypt_with_auth_data(self, mock_get_conn):
mock_get_conn.return_value.encrypt.return_value = RESPONSE
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,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual(PLAINTEXT_b64, result)
@mock.patch(
"airflow.providers.google.cloud.hooks.kms.CloudKMSHook.get_conn")
def test_decrypt(self, mock_get_conn):
mock_get_conn.return_value.decrypt.return_value = RESPONSE
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,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual(PLAINTEXT, result)
@mock.patch(
"airflow.providers.google.cloud.hooks.kms.CloudKMSHook.get_conn")
def test_decrypt_with_auth_data(self, mock_get_conn):
mock_get_conn.return_value.decrypt.return_value = RESPONSE
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,
retry=None,
timeout=None,
metadata=None,
)
self.assertEqual(PLAINTEXT, result)