[Storage] Add argument to perf tests to use client-side encryption (#…

…24978)

sdk/storage/azure-storage-blob/tests/perfstress_tests/README.md

@@ -54,6 +54,7 @@ The options are available for all Blob perf tests:
 - `--max-put-size` Maximum size of data uploading in single HTTP PUT. Default is 64\*1024\*1024.
 - `--max-block-size` Maximum size of data in a block within a blob. Defaults to 4\*1024\*1024.
 - `--buffer-threshold` Minimum block size to prevent full block buffering. Defaults to 4\*1024\*1024+1.
+- `--client-encryption` The version of client-side encryption to use. Leave out for no encryption.
 
 #### List Blobs command line options
 This option is only available to the List Blobs test (T1 and T2).

sdk/storage/azure-storage-blob/tests/perfstress_tests/_test_base.py

@@ -3,15 +3,14 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-import os
 import uuid
 
 from azure_devtools.perfstress_tests import PerfStressTest
-
-from azure.core.exceptions import ResourceNotFoundError
 from azure.storage.blob import BlobServiceClient as SyncBlobServiceClient
 from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
 
+from .key_wrapper import KeyWrapper
+
 
 class _ServiceTest(PerfStressTest):
     service_client = None
@@ -25,13 +24,18 @@ def __init__(self, arguments):
         self._client_kwargs['max_single_put_size'] = self.args.max_put_size
         self._client_kwargs['max_block_size'] = self.args.max_block_size
         self._client_kwargs['min_large_block_upload_threshold'] = self.args.buffer_threshold
+        if self.args.client_encryption:
+            self.key_encryption_key = KeyWrapper()
+            self._client_kwargs['require_encryption'] = True
+            self._client_kwargs['key_encryption_key'] = self.key_encryption_key
+            self._client_kwargs['encryption_version'] = self.args.client_encryption
         # self._client_kwargs['api_version'] = '2019-02-02'  # Used only for comparison with T1 legacy tests
 
         if not _ServiceTest.service_client or self.args.no_client_share:
             _ServiceTest.service_client = SyncBlobServiceClient.from_connection_string(conn_str=connection_string, **self._client_kwargs)
             _ServiceTest.async_service_client = AsyncBlobServiceClient.from_connection_string(conn_str=connection_string, **self._client_kwargs)
         self.service_client = _ServiceTest.service_client
-        self.async_service_client =_ServiceTest.async_service_client
+        self.async_service_client = _ServiceTest.async_service_client
 
     async def close(self):
         await self.async_service_client.close()
@@ -43,6 +47,7 @@ def add_arguments(parser):
         parser.add_argument('--max-put-size', nargs='?', type=int, help='Maximum size of data uploading in single HTTP PUT. Defaults to 64*1024*1024', default=64*1024*1024)
         parser.add_argument('--max-block-size', nargs='?', type=int, help='Maximum size of data in a block within a blob. Defaults to 4*1024*1024', default=4*1024*1024)
         parser.add_argument('--buffer-threshold', nargs='?', type=int, help='Minimum block size to prevent full block buffering. Defaults to 4*1024*1024+1', default=4*1024*1024+1)
+        parser.add_argument('--client-encryption', nargs='?', type=str, help='The version of client-side encryption to use. Leave out for no encryption.', default=None)
         parser.add_argument('--max-concurrency', nargs='?', type=int, help='Maximum number of concurrent threads used for data transfer. Defaults to 1', default=1)
         parser.add_argument('-s', '--size', nargs='?', type=int, help='Size of data to transfer.  Default is 10240.', default=10240)
         parser.add_argument('--no-client-share', action='store_true', help='Create one ServiceClient per test instance.  Default is to share a single ServiceClient.', default=False)

sdk/storage/azure-storage-blob/tests/perfstress_tests/key_wrapper.py

@@ -0,0 +1,34 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import os
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.keywrap import aes_key_wrap, aes_key_unwrap
+
+
+class KeyWrapper:
+    def __init__(self, kid='local:key1'):
+        self.kek = os.urandom(32)
+        self.backend = default_backend()
+        self.kid = kid
+
+    def wrap_key(self, key, algorithm='A256KW'):
+        if algorithm == 'A256KW':
+            return aes_key_wrap(self.kek, key, self.backend)
+
+        raise ValueError("Unknown key wrap algorithm.")
+
+    def unwrap_key(self, key, algorithm):
+        if algorithm == 'A256KW':
+            return aes_key_unwrap(self.kek, key, self.backend)
+
+        raise ValueError("Unknown key wrap algorithm.")
+
+    def get_key_wrap_algorithm(self):
+        return 'A256KW'
+
+    def get_kid(self):
+        return self.kid