From 05aec3dd91394cbc6517a5dbab7a91bfffba3446 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?McCoy=20Pati=C3=B1o?=
 <39780829+mccoyp@users.noreply.github.com>
Date: Thu, 9 Sep 2021 18:41:31 -0700
Subject: [PATCH] Redact account keys in recordings (#20537)

---
 ...st_mgmt_batch.test_mgmt_batch_account.yaml |  4 +--
 .../azure-mgmt-batch/tests/test_mgmt_batch.py | 35 +++++++++++++++++++
 2 files changed, 37 insertions(+), 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
index af80a5165f48..a9c6fc3d27ee 100644
--- 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
index b2ee86ecd5c4..64a66f5aa55b 100644
--- 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(