[Test Proxy] Add registration methods for all sanitizers

tools/azure-sdk-tools/devtools_testutils/__init__.py

@@ -1,5 +1,5 @@
 from .mgmt_testcase import AzureMgmtTestCase, AzureMgmtPreparer
-from .azure_recorded_testcase import add_sanitizer, AzureRecordedTestCase
+from .azure_recorded_testcase import AzureRecordedTestCase
 from .azure_testcase import AzureTestCase, is_live, get_region_override
 from .resource_testcase import (
     FakeResource,
@@ -16,12 +16,30 @@
 from .keyvault_preparer import KeyVaultPreparer
 from .powershell_preparer import PowerShellPreparer
 from .proxy_testcase import RecordedByProxy
-from .enums import ProxyRecordingSanitizer
+from .sanitizers import (
+    add_body_key_sanitizer,
+    add_body_regex_sanitizer,
+    add_continuation_sanitizer,
+    add_general_regex_sanitizer,
+    add_header_regex_sanitizer,
+    add_oauth_response_sanitizer,
+    add_remove_header_sanitizer,
+    add_request_subscription_id_sanitizer,
+    add_uri_regex_sanitizer,
+)
 from .helpers import ResponseCallback, RetryCounter
 from .fake_credential import FakeTokenCredential
 
 __all__ = [
-    "add_sanitizer",
+    "add_body_key_sanitizer",
+    "add_body_regex_sanitizer",
+    "add_continuation_sanitizer",
+    "add_general_regex_sanitizer",
+    "add_header_regex_sanitizer",
+    "add_oauth_response_sanitizer",
+    "add_remove_header_sanitizer",
+    "add_request_subscription_id_sanitizer",
+    "add_uri_regex_sanitizer",
     "AzureMgmtTestCase",
     "AzureMgmtPreparer",
     "AzureRecordedTestCase",
@@ -38,7 +56,6 @@
     "RandomNameResourceGroupPreparer",
     "CachedResourceGroupPreparer",
     "PowerShellPreparer",
-    "ProxyRecordingSanitizer",
     "RecordedByProxy",
     "ResponseCallback",
     "RetryCounter",

tools/azure-sdk-tools/devtools_testutils/azure_recorded_testcase.py

@@ -7,7 +7,6 @@
 import logging
 import os
 import os.path
-import requests
 import six
 import sys
 import time
@@ -21,8 +20,6 @@
 
 from . import mgmt_settings_fake as fake_settings
 from .azure_testcase import _is_autorest_v3, get_resource_name, get_qualified_method_name
-from .config import PROXY_URL
-from .enums import ProxyRecordingSanitizer
 
 try:
     # Try to import the AsyncFakeCredential, if we cannot assume it is Python 2
@@ -37,38 +34,6 @@
 load_dotenv(find_dotenv())
 
 
-def add_sanitizer(sanitizer, **kwargs):
-    # type: (ProxyRecordingSanitizer, **Any) -> None
-    """Registers a sanitizer, matcher, or transform with the test proxy.
-
-    :param sanitizer: The name of the sanitizer, matcher, or transform you want to add.
-    :type sanitizer: ProxyRecordingSanitizer or str
-
-    :keyword str value: The substitution value.
-    :keyword str regex: A regex for a sanitizer. Can be defined as a simple regex, or if a ``group_for_replace`` is
-        provided, a substitution operation.
-    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
-        a simple replacement operation.
-    """
-    request_args = {}
-    request_args["value"] = kwargs.get("value") or "fakevalue"
-    request_args["regex"] = (
-        kwargs.get("regex") or "(?<=\\/\\/)[a-z]+(?=(?:|-secondary)\\.(?:table|blob|queue)\\.core\\.windows\\.net)"
-    )
-    request_args["group_for_replace"] = kwargs.get("group_for_replace")
-
-    if sanitizer == ProxyRecordingSanitizer.URI:
-        requests.post(
-            "{}/Admin/AddSanitizer".format(PROXY_URL),
-            headers={"x-abstraction-identifier": ProxyRecordingSanitizer.URI.value},
-            json={
-                "regex": request_args["regex"],
-                "value": request_args["value"],
-                "groupForReplace": request_args["group_for_replace"],
-            },
-        )
-
-
 def is_live():
     """A module version of is_live, that could be used in pytest marker."""
     if not hasattr(is_live, "_cache"):

tools/azure-sdk-tools/devtools_testutils/sanitizers.py

@@ -0,0 +1,199 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+import requests
+from typing import TYPE_CHECKING
+
+from .config import PROXY_URL
+
+if TYPE_CHECKING:
+    from typing import Any, Dict
+
+
+def add_body_key_sanitizer(**kwargs):
+    # type: (**Any) -> None
+    """Registers a sanitizer that offers regex update of a specific JTokenPath within a returned body.
+
+    For example, "TableName" within a json response body having its value replaced by whatever substitution is offered.
+
+    :keyword str json_path: The SelectToken path (which could possibly match multiple entries) that will be used to
+        select JTokens for value replacement.
+    :keyword str value: The substitution value.
+    :keyword str regex: A regex. Can be defined as a simple regex replace OR if groupForReplace is set, a subsitution
+        operation. Defaults to replacing the entire string.
+    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
+        a simple replacement operation.
+    """
+
+    request_args = {}
+    request_args["jsonPath"] = kwargs.get("json_path")
+    request_args["value"] = kwargs.get("value")
+    request_args["regex"] = kwargs.get("regex")
+    request_args["groupForReplace"] = kwargs.get("group_for_replace")
+
+    _send_request("BodyKeySanitizer", request_args)
+
+
+def add_body_regex_sanitizer(**kwargs):
+    # type: (**Any) -> None
+    """Registers a sanitizer that offers regex replace within a returned body.
+
+    Specifically, this means regex applying to the raw JSON. If you are attempting to simply replace a specific key, the
+    BodyKeySanitizer is probably the way to go.
+
+    :keyword str value: The substitution value.
+    :keyword str regex: A regex. Can be defined as a simple regex, or if a ``group_for_replace`` is provided, a
+        substitution operation.
+    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
+        a simple replacement operation.
+    """
+
+    request_args = {}
+    request_args["value"] = kwargs.get("value")
+    request_args["regex"] = kwargs.get("regex")
+    request_args["groupForReplace"] = kwargs.get("group_for_replace")
+
+    _send_request("BodyRegexSanitizer", request_args)
+
+
+def add_continuation_sanitizer(**kwargs):
+    # type: (**Any) -> None
+    """Registers a sanitizer that's used to anonymize private keys in response/request pairs.
+
+    For instance, a request hands back a "sessionId" that needs to be present in the next request. Supports "all further
+    requests get this key" as well as "single response/request pair". Defaults to maintaining same key for rest of
+    recording.
+
+    :keyword str key: The name of the header whos value will be replaced from response -> next request.
+    :keyword str method: The method by which the value of the targeted key will be replaced. Defaults to guid
+        replacement.
+    :keyword str reset_after_first: Do we need multiple pairs replaced? Or do we want to replace each value with the
+        same value?
+    """
+
+    request_args = {}
+    request_args["key"] = kwargs.get("key")
+    request_args["method"] = kwargs.get("method")
+    request_args["resetAfterFirst"] = kwargs.get("reset_after_first")
+
+    _send_request("ContinuationSanitizer", request_args)
+
+
+def add_general_regex_sanitizer(**kwargs):
+    # type: (**Any) -> None
+    """Registers a sanitizer that offers a general regex replace across request/response Body, Headers, and URI.
+
+    For the body, this means regex applying to the raw JSON.
+
+    :keyword str value: The substitution value.
+    :keyword str regex: A regex. Can be defined as a simple regex, or if a ``group_for_replace`` is provided, a
+        substitution operation.
+    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
+        a simple replacement operation.
+    """
+
+    request_args = {}
+    request_args["value"] = kwargs.get("value")
+    request_args["regex"] = kwargs.get("regex")
+    request_args["groupForReplace"] = kwargs.get("group_for_replace")
+
+    _send_request("GeneralRegexSanitizer", request_args)
+
+
+def add_header_regex_sanitizer(**kwargs):
+    # type: (**Any) -> None
+    """Registers a sanitizer that offers regex replace on returned headers.
+
+    Can be used for multiple purposes: 1) To replace a key with a specific value, do not set "regex" value. 2) To do a
+    simple regex replace operation, define arguments "key", "value", and "regex". 3) To do a targeted substitution of a
+    specific group, define all arguments "key", "value", and "regex".
+
+    :keyword str key: The name of the header we're operating against.
+    :keyword str value: The substitution or whole new header value, depending on "regex" setting.
+    :keyword str regex: A regex. Can be defined as a simple regex, or if a ``group_for_replace`` is provided, a
+        substitution operation.
+    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
+        a simple replacement operation.
+    """
+
+    request_args = {}
+    request_args["key"] = kwargs.get("key")
+    request_args["value"] = kwargs.get("value")
+    request_args["regex"] = kwargs.get("regex")
+    request_args["groupForReplace"] = kwargs.get("group_for_replace")
+
+    _send_request("HeaderRegexSanitizer", request_args)
+
+
+def add_oauth_response_sanitizer():
+    # type: () -> None
+    """Registers a sanitizer that cleans out all request/response pairs that match an oauth regex in their URI."""
+
+    _send_request("OAuthResponseSanitizer", {})
+
+
+def add_remove_header_sanitizer(**kwargs):
+    # type: (**Any) -> None
+    """Registers a sanitizer that removes specified headers before saving a recording.
+
+    :keyword str headers: A comma separated list. Should look like "Location, Transfer-Encoding" or something along
+        those lines. Don't worry about whitespace between the commas separating each key. They will be ignored.
+    """
+
+    request_args = {}
+    request_args["headersForRemoval"] = kwargs.get("headers")
+
+    _send_request("RemoveHeaderSanitizer", request_args)
+
+
+def add_request_subscription_id_sanitizer(**kwargs):
+    # type: (**Any) -> None
+    """Registers a sanitizer that replaces subscription IDs in requests.
+
+    Subscription IDs are replaced with "00000000-0000-0000-0000-000000000000" by default.
+
+    :keyword str value: The fake subscriptionId that will be placed where the real one is in the real request.
+    """
+
+    request_args = {}
+    request_args["value"] = kwargs.get("value") or "fakevalue"
+
+    _send_request("ReplaceRequestSubscriptionId", request_args)
+
+
+def add_uri_regex_sanitizer(**kwargs):
+    # type: (**Any) -> None
+    """Registers a sanitizer for cleaning URIs via regex.
+
+    :keyword str value: The substitution value.
+    :keyword str regex: A regex. Can be defined as a simple regex, or if a ``group_for_replace`` is provided, a
+        substitution operation.
+    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
+        a simple replacement operation.
+    """
+
+    request_args = {}
+    request_args["value"] = kwargs.get("value") or "fakevalue"
+    request_args["regex"] = (
+        kwargs.get("regex") or "(?<=\\/\\/)[a-z]+(?=(?:|-secondary)\\.(?:table|blob|queue)\\.core\\.windows\\.net)"
+    )
+    request_args["groupForReplace"] = kwargs.get("group_for_replace")
+
+    _send_request("UriRegexSanitizer", request_args)
+
+
+def _send_request(sanitizer, parameters):
+    # type: (str, Dict) -> None
+    """Send a POST request to the test proxy endpoint to register the specified sanitizer.
+
+    :param str sanitizer: The name of the sanitizer to add.
+    :param dict parameters: The sanitizer constructor parameters, as a dictionary.
+    """
+
+    requests.post(
+        "{}/Admin/AddSanitizer".format(PROXY_URL),
+        headers={"x-abstraction-identifier": sanitizer},
+        json=parameters
+    )