# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
# pylint: disable=too-many-public-methods, unnecessary-lambda-assignment

import os
import unittest
from base64 import b64decode, b64encode
from json import dumps, loads
from unittest import mock

import pytest

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

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


# ------------------------------------------------------------------------------
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:
            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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret)
        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.secret,
            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.secret,
            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.secret,
            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.secret,
            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.secret,
            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.secret,
            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:
            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.secret,
            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.secret,
            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.secret,
            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()
