# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- import os import unittest from base64 import b64decode, b64encode from json import dumps, loads from unittest import mock import pytest from azure.core.exceptions import ResourceExistsError, HttpResponseError from azure.storage.queue import ( BinaryBase64DecodePolicy, BinaryBase64EncodePolicy, QueueServiceClient, VERSION, ) from azure.storage.queue._encryption import ( _dict_to_encryption_data, _EncryptionAgent, _EncryptionData, _ERROR_OBJECT_INVALID, _GCM_NONCE_LENGTH, _GCM_TAG_LENGTH, _validate_and_unwrap_cek, _WrappedContentKey, ) from azure.storage.queue._shared import decode_base64_to_bytes from cryptography.hazmat import backends from cryptography.hazmat.primitives.ciphers import Cipher from cryptography.hazmat.primitives.ciphers.aead import AESGCM from cryptography.hazmat.primitives.ciphers.algorithms import AES from cryptography.hazmat.primitives.ciphers.modes import CBC from cryptography.hazmat.primitives.padding import PKCS7 from devtools_testutils import recorded_by_proxy from devtools_testutils.storage import StorageRecordedTestCase from encryption_test_helper import KeyResolver, KeyWrapper, mock_urandom, RSAKeyWrapper from settings.testcase import QueuePreparer # ------------------------------------------------------------------------------ TEST_QUEUE_PREFIX = "encryptionqueue" # ------------------------------------------------------------------------------ def _decode_base64_to_bytes(data): if isinstance(data, str): data = data.encode("utf-8") return b64decode(data) @mock.patch("os.urandom", mock_urandom) class TestStorageQueueEncryption(StorageRecordedTestCase): # --Helpers----------------------------------------------------------------- def _get_queue_reference(self, qsc, prefix=TEST_QUEUE_PREFIX, **kwargs): queue_name = self.get_resource_name(prefix) queue = qsc.get_queue_client(queue_name, **kwargs) return queue def _create_queue(self, qsc, prefix=TEST_QUEUE_PREFIX, **kwargs): queue = self._get_queue_reference(qsc, prefix, **kwargs) try: created = queue.create_queue() except ResourceExistsError: pass return queue # -------------------------------------------------------------------------- @QueuePreparer() @recorded_by_proxy def test_get_messages_encrypted_kek(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) qsc.key_encryption_key = KeyWrapper("key1") queue = self._create_queue(qsc) queue.send_message("encrypted_message_2") # Act li = next(queue.receive_messages()) # Assert assert li.content == "encrypted_message_2" @QueuePreparer() @recorded_by_proxy def test_get_messages_encrypted_resolver(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) qsc.key_encryption_key = KeyWrapper("key1") queue = self._create_queue(qsc) queue.send_message("encrypted_message_2") key_resolver = KeyResolver() key_resolver.put_key(qsc.key_encryption_key) queue.key_resolver_function = key_resolver.resolve_key queue.key_encryption_key = None # Ensure that the resolver is used # Act li = next(queue.receive_messages()) # Assert assert li.content == "encrypted_message_2" @QueuePreparer() @recorded_by_proxy def test_peek_messages_encrypted_kek(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) qsc.key_encryption_key = KeyWrapper("key1") queue = self._create_queue(qsc) queue.send_message("encrypted_message_3") # Act li = queue.peek_messages() # Assert assert li[0].content == "encrypted_message_3" @QueuePreparer() @recorded_by_proxy def test_peek_messages_encrypted_resolver(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) qsc.key_encryption_key = KeyWrapper("key1") queue = self._create_queue(qsc) queue.send_message("encrypted_message_4") key_resolver = KeyResolver() key_resolver.put_key(qsc.key_encryption_key) queue.key_resolver_function = key_resolver.resolve_key queue.key_encryption_key = None # Ensure that the resolver is used # Act li = queue.peek_messages() # Assert assert li[0].content == "encrypted_message_4" @pytest.mark.live_test_only @QueuePreparer() def test_peek_messages_encrypted_kek_RSA(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # We can only generate random RSA keys, so this must be run live or # the playback test will fail due to a change in kek values. # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) qsc.key_encryption_key = RSAKeyWrapper("key2") queue = self._create_queue(qsc) queue.send_message("encrypted_message_3") # Act li = queue.peek_messages() # Assert assert li[0].content == "encrypted_message_3" @QueuePreparer() @recorded_by_proxy def test_update_encrypted_message(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) queue.key_encryption_key = KeyWrapper("key1") queue.send_message("Update Me") messages = queue.receive_messages() list_result1 = next(messages) list_result1.content = "Updated" # Act message = queue.update_message(list_result1) list_result2 = next(messages) # Assert assert "Updated" == list_result2.content @QueuePreparer() @recorded_by_proxy def test_update_encrypted_binary_message(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue( qsc, message_encode_policy=BinaryBase64EncodePolicy(), message_decode_policy=BinaryBase64DecodePolicy() ) queue.key_encryption_key = KeyWrapper("key1") binary_message = self.get_random_bytes(100) queue.send_message(binary_message) messages = [] for m in queue.receive_messages(): messages.append(m) list_result1 = messages[0] # Act binary_message = self.get_random_bytes(100) list_result1.content = binary_message queue.update_message(list_result1) for m in queue.receive_messages(): messages.append(m) list_result2 = messages[0] # Assert assert binary_message == list_result2.content @QueuePreparer() @recorded_by_proxy def test_update_encrypted_raw_text_message(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc, message_encode_policy=None, message_decode_policy=None) queue.key_encryption_key = KeyWrapper("key1") raw_text = "Update Me" queue.send_message(raw_text) messages = queue.receive_messages() list_result1 = next(messages) # Act raw_text = "Updated" list_result1.content = raw_text queue.update_message(list_result1) list_result2 = next(messages) # Assert assert raw_text == list_result2.content @QueuePreparer() @recorded_by_proxy def test_update_encrypted_json_message(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc, message_encode_policy=None, message_decode_policy=None) queue.key_encryption_key = KeyWrapper("key1") message_dict = {"val1": 1, "val2": "2"} json_text = dumps(message_dict) queue.send_message(json_text) messages = queue.receive_messages() list_result1 = next(messages) # Act message_dict["val1"] = 0 message_dict["val2"] = "updated" json_text = dumps(message_dict) list_result1.content = json_text queue.update_message(list_result1) list_result2 = next(messages) # Assert assert message_dict == loads(list_result2.content) @QueuePreparer() @recorded_by_proxy def test_invalid_value_kek_wrap(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) queue.key_encryption_key = KeyWrapper("key1") queue.key_encryption_key.get_kid = None with pytest.raises(AttributeError) as e: queue.send_message("message") assert str(e.value.args[0]), _ERROR_OBJECT_INVALID.format("key encryption key" == "get_kid") queue.key_encryption_key = KeyWrapper("key1") queue.key_encryption_key.get_kid = None with pytest.raises(AttributeError): queue.send_message("message") queue.key_encryption_key = KeyWrapper("key1") queue.key_encryption_key.wrap_key = None with pytest.raises(AttributeError): queue.send_message("message") @QueuePreparer() @recorded_by_proxy def test_missing_attribute_kek_wrap(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) valid_key = KeyWrapper("key1") # Act invalid_key_1 = lambda: None # functions are objects, so this effectively creates an empty object invalid_key_1.get_key_wrap_algorithm = valid_key.get_key_wrap_algorithm invalid_key_1.get_kid = valid_key.get_kid # No attribute wrap_key queue.key_encryption_key = invalid_key_1 with pytest.raises(AttributeError): queue.send_message("message") invalid_key_2 = lambda: None # functions are objects, so this effectively creates an empty object invalid_key_2.wrap_key = valid_key.wrap_key invalid_key_2.get_kid = valid_key.get_kid # No attribute get_key_wrap_algorithm queue.key_encryption_key = invalid_key_2 with pytest.raises(AttributeError): queue.send_message("message") invalid_key_3 = lambda: None # functions are objects, so this effectively creates an empty object invalid_key_3.get_key_wrap_algorithm = valid_key.get_key_wrap_algorithm invalid_key_3.wrap_key = valid_key.wrap_key # No attribute get_kid queue.key_encryption_key = invalid_key_3 with pytest.raises(AttributeError): queue.send_message("message") @QueuePreparer() @recorded_by_proxy def test_invalid_value_kek_unwrap(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) queue.key_encryption_key = KeyWrapper("key1") queue.send_message("message") # Act queue.key_encryption_key.unwrap_key = None with pytest.raises(HttpResponseError): queue.peek_messages() queue.key_encryption_key.get_kid = None with pytest.raises(HttpResponseError): queue.peek_messages() @QueuePreparer() @recorded_by_proxy def test_missing_attribute_kek_unwrap(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) queue.key_encryption_key = KeyWrapper("key1") queue.send_message("message") # Act valid_key = KeyWrapper("key1") invalid_key_1 = lambda: None # functions are objects, so this effectively creates an empty object invalid_key_1.unwrap_key = valid_key.unwrap_key # No attribute get_kid queue.key_encryption_key = invalid_key_1 with pytest.raises(HttpResponseError) as e: queue.peek_messages() assert "Decryption failed." in str(e.value.args[0]) invalid_key_2 = lambda: None # functions are objects, so this effectively creates an empty object invalid_key_2.get_kid = valid_key.get_kid # No attribute unwrap_key queue.key_encryption_key = invalid_key_2 with pytest.raises(HttpResponseError): queue.peek_messages() @QueuePreparer() @recorded_by_proxy def test_validate_encryption(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) kek = KeyWrapper("key1") queue.key_encryption_key = kek queue.send_message("message") # Act queue.key_encryption_key = None # Message will not be decrypted li = queue.peek_messages() message = li[0].content message = loads(message) encryption_data = message["EncryptionData"] wrapped_content_key = encryption_data["WrappedContentKey"] wrapped_content_key = _WrappedContentKey( wrapped_content_key["Algorithm"], b64decode(wrapped_content_key["EncryptedKey"].encode(encoding="utf-8")), wrapped_content_key["KeyId"], ) encryption_agent = encryption_data["EncryptionAgent"] encryption_agent = _EncryptionAgent(encryption_agent["EncryptionAlgorithm"], encryption_agent["Protocol"]) encryption_data = _EncryptionData( b64decode(encryption_data["ContentEncryptionIV"].encode(encoding="utf-8")), None, encryption_agent, wrapped_content_key, {"EncryptionLibrary": VERSION}, ) message = message["EncryptedMessageContents"] content_encryption_key = kek.unwrap_key( encryption_data.wrapped_content_key.encrypted_key, encryption_data.wrapped_content_key.algorithm ) # Create decryption cipher backend = backends.default_backend() algorithm = AES(content_encryption_key) mode = CBC(encryption_data.content_encryption_IV) cipher = Cipher(algorithm, mode, backend) # decode and decrypt data decrypted_data = decode_base64_to_bytes(message) decryptor = cipher.decryptor() decrypted_data = decryptor.update(decrypted_data) + decryptor.finalize() # unpad data unpadder = PKCS7(128).unpadder() decrypted_data = unpadder.update(decrypted_data) + unpadder.finalize() decrypted_data = decrypted_data.decode(encoding="utf-8") # Assert assert decrypted_data == "message" @QueuePreparer() @recorded_by_proxy def test_put_with_strict_mode(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) kek = KeyWrapper("key1") queue.key_encryption_key = kek queue.require_encryption = True queue.send_message("message") queue.key_encryption_key = None # Assert with pytest.raises(ValueError) as e: queue.send_message("message") assert str(e.value.args[0]) == "Encryption required but no key was provided." @QueuePreparer() @recorded_by_proxy def test_get_with_strict_mode(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) queue.send_message("message") queue.require_encryption = True queue.key_encryption_key = KeyWrapper("key1") with pytest.raises(ValueError) as e: next(queue.receive_messages()) assert "Message was either not encrypted or metadata was incorrect." in str(e.value.args[0]) @QueuePreparer() @recorded_by_proxy def test_encryption_add_encrypted_64k_message(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) message = "a" * 1024 * 64 # Act queue.send_message(message) # Assert queue.key_encryption_key = KeyWrapper("key1") with pytest.raises(HttpResponseError): queue.send_message(message) @QueuePreparer() @recorded_by_proxy def test_encryption_nonmatching_kid(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient(self.account_url(storage_account_name, "queue"), storage_account_key) queue = self._create_queue(qsc) queue.key_encryption_key = KeyWrapper("key1") queue.send_message("message") # Act queue.key_encryption_key.kid = "Invalid" # Assert with pytest.raises(HttpResponseError) as e: next(queue.receive_messages()) assert "Decryption failed." in str(e.value.args[0]) @QueuePreparer() @recorded_by_proxy def test_get_message_encrypted_kek_v2(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient( self.account_url(storage_account_name, "queue"), storage_account_key, require_encryption=True, encryption_version="2.0", key_encryption_key=KeyWrapper("key1"), ) queue = self._create_queue(qsc) content = "Hello World Encrypted!" # Act queue.send_message(content) message = queue.receive_message() # Assert assert content == message.content @QueuePreparer() @recorded_by_proxy def test_get_message_encrypted_resolver_v2(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient( self.account_url(storage_account_name, "queue"), storage_account_key, require_encryption=True, encryption_version="2.0", key_encryption_key=KeyWrapper("key1"), ) key_resolver = KeyResolver() key_resolver.put_key(qsc.key_encryption_key) queue = self._create_queue(qsc) content = "Hello World Encrypted!" # Act queue.send_message(content) queue.key_resolver_function = key_resolver.resolve_key queue.key_encryption_key = None # Ensure that the resolver is used message = queue.receive_message() # Assert assert content == message.content @pytest.mark.live_test_only @QueuePreparer() def test_get_message_encrypted_kek_RSA_v2(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # We can only generate random RSA keys, so this must be run live or # the playback test will fail due to a change in kek values. # Arrange qsc = QueueServiceClient( self.account_url(storage_account_name, "queue"), storage_account_key, require_encryption=True, encryption_version="2.0", key_encryption_key=RSAKeyWrapper("key2"), ) queue = self._create_queue(qsc) content = "Hello World Encrypted!" # Act queue.send_message(content) message = queue.receive_message() # Assert assert content == message.content @QueuePreparer() @recorded_by_proxy def test_update_encrypted_message_v2(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient( self.account_url(storage_account_name, "queue"), storage_account_key, requires_encryption=True, encryption_version="2.0", key_encryption_key=KeyWrapper("key1"), ) queue = self._create_queue(qsc) queue.send_message("Update Me") message = queue.receive_message() message.content = "Updated" # Act queue.update_message(message) message = queue.receive_message() # Assert assert "Updated" == message.content @QueuePreparer() @recorded_by_proxy def test_update_encrypted_binary_message_v2(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange qsc = QueueServiceClient( self.account_url(storage_account_name, "queue"), storage_account_key, requires_encryption=True, encryption_version="2.0", key_encryption_key=KeyWrapper("key1"), ) queue = self._create_queue( qsc, message_encode_policy=BinaryBase64EncodePolicy(), message_decode_policy=BinaryBase64DecodePolicy() ) queue.key_encryption_key = KeyWrapper("key1") queue.send_message(b"Update Me") message = queue.receive_message() message.content = b"Updated" # Act queue.update_message(message) message = queue.receive_message() # Assert assert b"Updated" == message.content @QueuePreparer() @recorded_by_proxy def test_encryption_v2_v1_downgrade(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange kek = KeyWrapper("key1") qsc = QueueServiceClient( self.account_url(storage_account_name, "queue"), storage_account_key, requires_encryption=True, encryption_version="2.0", key_encryption_key=kek, ) queue = self._create_queue(qsc) queue.send_message("Hello World Encrypted!") queue.require_encryption = False queue.key_encryption_key = None # Message will not be decrypted message = queue.receive_message() content = loads(message.content) # Modify metadata to look like V1 encryption_data = content["EncryptionData"] encryption_data["EncryptionAgent"]["Protocol"] = "1.0" encryption_data["EncryptionAgent"]["EncryptionAlgorithm"] = "AES_CBC_256" iv = b64encode(os.urandom(16)) encryption_data["ContentEncryptionIV"] = iv.decode("utf-8") content["EncryptionData"] = encryption_data message.content = dumps(content) # Act / Assert # Update without encryption queue.update_message(message) # Re-enable encryption for receive queue.require_encryption = True queue.key_encryption_key = kek with pytest.raises(HttpResponseError) as e: new_message = queue.receive_message() assert "Decryption failed." in str(e.value.args[0]) @QueuePreparer() @recorded_by_proxy def test_validate_encryption_v2(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") # Arrange kek = KeyWrapper("key1") qsc = QueueServiceClient( self.account_url(storage_account_name, "queue"), storage_account_key, require_encryption=True, encryption_version="2.0", key_encryption_key=kek, ) queue = self._create_queue(qsc) content = "Hello World Encrypted!" queue.send_message(content) # Act queue.require_encryption = False queue.key_encryption_key = None # Message will not be decrypted message = queue.receive_message().content message = loads(message) encryption_data = _dict_to_encryption_data(message["EncryptionData"]) encryption_agent = encryption_data.encryption_agent assert "2.0" == encryption_agent.protocol assert "AES_GCM_256" == encryption_agent.encryption_algorithm encrypted_region_info = encryption_data.encrypted_region_info assert _GCM_NONCE_LENGTH == encrypted_region_info.nonce_length assert _GCM_TAG_LENGTH == encrypted_region_info.tag_length content_encryption_key = _validate_and_unwrap_cek(encryption_data, kek, None) nonce_length = encrypted_region_info.nonce_length message = message["EncryptedMessageContents"] message = decode_base64_to_bytes(message) # First bytes are the nonce nonce = message[:nonce_length] ciphertext_with_tag = message[nonce_length:] aesgcm = AESGCM(content_encryption_key) decrypted_data = aesgcm.decrypt(nonce, ciphertext_with_tag, None) decrypted_data = decrypted_data.decode(encoding="utf-8") # Assert assert content == decrypted_data @QueuePreparer() @recorded_by_proxy def test_encryption_user_agent(self, **kwargs): storage_account_name = kwargs.pop("storage_account_name") storage_account_key = kwargs.pop("storage_account_key") app_id = "TestAppId" content = "Hello World Encrypted!" kek = KeyWrapper("key1") def assert_user_agent(request): start = f"{app_id} azstorage-clientsideencryption/2.0 " assert request.http_request.headers["User-Agent"].startswith(start) # Test method level keyword qsc = QueueServiceClient( self.account_url(storage_account_name, "queue"), storage_account_key, require_encryption=True, encryption_version="2.0", key_encryption_key=kek, ) queue = self._create_queue(qsc) queue.send_message(content, raw_request_hook=assert_user_agent, user_agent=app_id) # Test client constructor level keyword qsc = QueueServiceClient( self.account_url(storage_account_name, "queue"), storage_account_key, require_encryption=True, encryption_version="2.0", key_encryption_key=kek, user_agent=app_id, ) queue = self._get_queue_reference(qsc) queue.send_message(content, raw_request_hook=assert_user_agent) # ------------------------------------------------------------------------------ if __name__ == "__main__": unittest.main()