Redact account keys in recordings (#20537)

Azure · Sep 10, 2021 · 05aec3d · 05aec3d
1 parent 6a9cb86
commit 05aec3d
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 2 deletions.
diff --git a/sdk/batch/azure-mgmt-batch/tests/recordings/test_mgmt_batch.test_mgmt_batch_account.yaml b/sdk/batch/azure-mgmt-batch/tests/recordings/test_mgmt_batch.test_mgmt_batch_account.yaml
@@ -200,7 +200,7 @@ interactions:
     uri: https://centraluseuap.management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/test_mgmt_batch_test_mgmt_batch_account3e1b0fe5/providers/Microsoft.Batch/batchAccounts/batch3e1b0fe5/listKeys?api-version=2021-06-01
   response:
     body:
-      string: '{"accountName":"batch3e1b0fe5","primary":"3UQ9ry1mRmgftC37/IOylMEnaC713zLTXoMqp/zBQZ1ANY8eLsv1j5lkvN3PnaSevqoKjtfjKFyJ5Vsc6SGA0w==","secondary":"6dM/Myi6VRmOwbqgMjcIv4lSS7SQvlSTCmQX3RiwLvbivKU9oFi5zgdx7oNtOATEbB9rYO8oDkYVwn8PJLaTcg=="}'
+      string: '{"accountName":"batch3e1b0fe5","primary":"redacted6f7d7a","secondary":"redacted59d978"}'
     headers:
       cache-control:
       - no-cache
@@ -248,7 +248,7 @@ interactions:
     uri: https://centraluseuap.management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/test_mgmt_batch_test_mgmt_batch_account3e1b0fe5/providers/Microsoft.Batch/batchAccounts/batch3e1b0fe5/regenerateKeys?api-version=2021-06-01
   response:
     body:
-      string: '{"accountName":"batch3e1b0fe5","primary":"3UQ9ry1mRmgftC37/IOylMEnaC713zLTXoMqp/zBQZ1ANY8eLsv1j5lkvN3PnaSevqoKjtfjKFyJ5Vsc6SGA0w==","secondary":"Q3pqv2ncSAxxnhTR14lumWnq9GRUVvy8exfqF2q5x6SZYGpregob+HI5eehGuFusbCaHLdzdzr3ZqBKXY3Qtyw=="}'
+      string: '{"accountName":"batch3e1b0fe5","primary":"redacted6f7d7a","secondary":"redactedebdb55"}'
     headers:
       cache-control:
       - no-cache

diff --git a/sdk/batch/azure-mgmt-batch/tests/test_mgmt_batch.py b/sdk/batch/azure-mgmt-batch/tests/test_mgmt_batch.py
@@ -5,18 +5,23 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 #--------------------------------------------------------------------------
+import binascii
+import hashlib
 import io
+import json
 import logging
 import time
 import unittest
 
 import requests
+import six
 
 import azure.mgmt.batch
 from azure.mgmt.batch import models
 import azure.mgmt.network.models
 from mgmt_batch_preparers import KeyVaultPreparer, SimpleBatchPreparer
 
+from azure_devtools.scenario_tests.recording_processors import GeneralNameReplacer, RecordingProcessor
 from devtools_testutils import (
     AzureMgmtTestCase,
     ResourceGroupPreparer,
@@ -32,10 +37,40 @@
 EXPECTED_DEDICATED_CORE_QUOTA = 500
 EXPECTED_LOW_PRIO_CORE_QUOTA = 500
 EXPECTED_POOL_QUOTA = 100
+SECRET_FIELDS = ["primary", "secondary"]
+
+
+def get_redacted_key(key):
+    redacted_value = "redacted"
+    digest = hashlib.sha256(six.ensure_binary(key)).digest()
+    redacted_value += six.ensure_str(binascii.hexlify(digest))[:6]
+    return redacted_value
+
+
+class RecordingRedactor(RecordingProcessor):
+    """Removes keys from test recordings"""
+
+    def process_response(self, response):
+        try:
+            body = json.loads(response["body"]["string"])
+        except (KeyError, ValueError):
+            return response
+
+        for field in body:
+            if field in SECRET_FIELDS:
+                body[field] = get_redacted_key(body[field])
+
+        response["body"]["string"] = json.dumps(body)
+        return response
 
 
 class MgmtBatchTest(AzureMgmtTestCase):
 
+    def __init__(self, *args, **kwargs):
+        scrubber = GeneralNameReplacer()
+        redactor = RecordingRedactor()
+        super(MgmtBatchTest, self).__init__(*args, recording_processors=[redactor, scrubber], **kwargs)
+
     def setUp(self):
         super(MgmtBatchTest, self).setUp()
         self.mgmt_batch_client = self.create_mgmt_client(