diff --git a/tools/azure-sdk-tools/devtools_testutils/__init__.py b/tools/azure-sdk-tools/devtools_testutils/__init__.py index c58ea74788c2..b5303641ae36 100644 --- a/tools/azure-sdk-tools/devtools_testutils/__init__.py +++ b/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", diff --git a/tools/azure-sdk-tools/devtools_testutils/azure_recorded_testcase.py b/tools/azure-sdk-tools/devtools_testutils/azure_recorded_testcase.py index 969d80fe57fe..c2a709d0dc91 100644 --- a/tools/azure-sdk-tools/devtools_testutils/azure_recorded_testcase.py +++ b/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"): diff --git a/tools/azure-sdk-tools/devtools_testutils/enums.py b/tools/azure-sdk-tools/devtools_testutils/enums.py deleted file mode 100644 index f1456dbc8a1f..000000000000 --- a/tools/azure-sdk-tools/devtools_testutils/enums.py +++ /dev/null @@ -1,11 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. See License.txt in the project root for -# license information. -# -------------------------------------------------------------------------- -from enum import Enum - -class ProxyRecordingSanitizer(str, Enum): - """General-purpose sanitizers for sanitizing test proxy recordings""" - - URI = "UriRegexSanitizer" diff --git a/tools/azure-sdk-tools/devtools_testutils/sanitizers.py b/tools/azure-sdk-tools/devtools_testutils/sanitizers.py new file mode 100644 index 000000000000..0922828fea31 --- /dev/null +++ b/tools/azure-sdk-tools/devtools_testutils/sanitizers.py @@ -0,0 +1,183 @@ +# ------------------------------------------------------------------------- +# 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 = _get_request_args(**kwargs) + _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 = _get_request_args(**kwargs) + _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 = _get_request_args(**kwargs) + _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 = _get_request_args(**kwargs) + _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 = _get_request_args(**kwargs) + _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 = _get_request_args(**kwargs) + _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 = _get_request_args(**kwargs) + _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 = _get_request_args(**kwargs) + _send_request("UriRegexSanitizer", request_args) + + +def _get_request_args(**kwargs): + # type: (**Any) -> Dict + """Returns a dictionary of sanitizer constructor headers""" + + request_args = {} + request_args["groupForReplace"] = kwargs.get("group_for_replace") + request_args["headersForRemoval"] = kwargs.get("headers") + request_args["jsonPath"] = kwargs.get("json_path") + request_args["key"] = kwargs.get("key") + request_args["method"] = kwargs.get("method") + request_args["regex"] = kwargs.get("regex") + request_args["resetAfterFirst"] = kwargs.get("reset_after_first") + request_args["value"] = kwargs.get("value") + return 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 + )