1# -------------------------------------------------------------------------- 
    2# 
    3# Copyright (c) Microsoft Corporation. All rights reserved. 
    4# 
    5# The MIT License (MIT) 
    6# 
    7# Permission is hereby granted, free of charge, to any person obtaining a copy 
    8# of this software and associated documentation files (the ""Software""), to 
    9# deal in the Software without restriction, including without limitation the 
    10# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 
    11# sell copies of the Software, and to permit persons to whom the Software is 
    12# furnished to do so, subject to the following conditions: 
    13# 
    14# The above copyright notice and this permission notice shall be included in 
    15# all copies or substantial portions of the Software. 
    16# 
    17# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
    18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
    19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
    20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
    21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 
    22# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 
    23# IN THE SOFTWARE. 
    24# 
    25# -------------------------------------------------------------------------- 
    26from typing import List, Optional, Any, TypeVar 
    27from azure.core.pipeline import PipelineRequest 
    28from azure.core.pipeline.transport import ( 
    29    HttpRequest as LegacyHttpRequest, 
    30    HttpResponse as LegacyHttpResponse, 
    31) 
    32from azure.core.rest import HttpRequest, HttpResponse 
    33from ._base import SansIOHTTPPolicy 
    34 
    35HTTPResponseType = TypeVar("HTTPResponseType", HttpResponse, LegacyHttpResponse) 
    36HTTPRequestType = TypeVar("HTTPRequestType", HttpRequest, LegacyHttpRequest) 
    37 
    38 
    39class SensitiveHeaderCleanupPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseType]): 
    40    """A simple policy that cleans up sensitive headers 
    41 
    42    :keyword list[str] blocked_redirect_headers: The headers to clean up when redirecting to another domain. 
    43    :keyword bool disable_redirect_cleanup: Opt out cleaning up sensitive headers when redirecting to another domain. 
    44    """ 
    45 
    46    DEFAULT_SENSITIVE_HEADERS = set( 
    47        [ 
    48            "Authorization", 
    49            "x-ms-authorization-auxiliary", 
    50        ] 
    51    ) 
    52 
    53    def __init__( 
    54        self,  # pylint: disable=unused-argument 
    55        *, 
    56        blocked_redirect_headers: Optional[List[str]] = None, 
    57        disable_redirect_cleanup: bool = False, 
    58        **kwargs: Any 
    59    ) -> None: 
    60        self._disable_redirect_cleanup = disable_redirect_cleanup 
    61        self._blocked_redirect_headers = ( 
    62            SensitiveHeaderCleanupPolicy.DEFAULT_SENSITIVE_HEADERS 
    63            if blocked_redirect_headers is None 
    64            else blocked_redirect_headers 
    65        ) 
    66 
    67    def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None: 
    68        """This is executed before sending the request to the next policy. 
    69 
    70        :param request: The PipelineRequest object. 
    71        :type request: ~azure.core.pipeline.PipelineRequest 
    72        """ 
    73        # "insecure_domain_change" is used to indicate that a redirect 
    74        # has occurred to a different domain. This tells the SensitiveHeaderCleanupPolicy 
    75        # to clean up sensitive headers. We need to remove it before sending the request 
    76        # to the transport layer. 
    77        insecure_domain_change = request.context.options.pop("insecure_domain_change", False) 
    78        if not self._disable_redirect_cleanup and insecure_domain_change: 
    79            for header in self._blocked_redirect_headers: 
    80                request.http_request.headers.pop(header, None)