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# -------------------------------------------------------------------------- 
    26import asyncio 
    27import json 
    28import logging 
    29import uuid 
    30from typing import Union 
    31 
    32from azure.core.pipeline import PipelineRequest, PipelineResponse 
    33from azure.core.pipeline.policies import AsyncHTTPPolicy 
    34from azure.core.pipeline.transport import ( 
    35    HttpRequest as LegacyHttpRequest, 
    36    AsyncHttpResponse as LegacyAsyncHttpResponse, 
    37) 
    38from azure.core.rest import HttpRequest, AsyncHttpResponse 
    39 
    40 
    41from ._base import _SansIOARMAutoResourceProviderRegistrationPolicy 
    42 
    43_LOGGER = logging.getLogger(__name__) 
    44 
    45HTTPRequestType = Union[LegacyHttpRequest, HttpRequest] 
    46AsyncHTTPResponseType = Union[LegacyAsyncHttpResponse, AsyncHttpResponse] 
    47PipelineResponseType = PipelineResponse[HTTPRequestType, AsyncHTTPResponseType] 
    48 
    49 
    50class AsyncARMAutoResourceProviderRegistrationPolicy( 
    51    _SansIOARMAutoResourceProviderRegistrationPolicy, AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType] 
    52):  # pylint: disable=name-too-long 
    53    """Auto register an ARM resource provider if not done yet.""" 
    54 
    55    async def send( 
    56        self, request: PipelineRequest[HTTPRequestType] 
    57    ) -> PipelineResponse[HTTPRequestType, AsyncHTTPResponseType]: 
    58        http_request = request.http_request 
    59        response = await self.next.send(request) 
    60        if response.http_response.status_code == 409: 
    61            rp_name = self._check_rp_not_registered_err(response) 
    62            if rp_name: 
    63                url_prefix = self._extract_subscription_url(http_request.url) 
    64                register_rp_status = await self._async_register_rp(request, url_prefix, rp_name) 
    65                if not register_rp_status: 
    66                    return response 
    67                # Change the 'x-ms-client-request-id' otherwise the Azure endpoint 
    68                # just returns the same 409 payload without looking at the actual query 
    69                if "x-ms-client-request-id" in http_request.headers: 
    70                    http_request.headers["x-ms-client-request-id"] = str(uuid.uuid4()) 
    71                response = await self.next.send(request) 
    72        return response 
    73 
    74    async def _async_register_rp( 
    75        self, initial_request: PipelineRequest[HTTPRequestType], url_prefix: str, rp_name: str 
    76    ) -> bool: 
    77        """Synchronously register the RP is paremeter. 
    78 
    79        Return False if we have a reason to believe this didn't work 
    80 
    81        :param initial_request: The initial request 
    82        :type initial_request: ~azure.core.pipeline.PipelineRequest 
    83        :param str url_prefix: The url prefix 
    84        :param str rp_name: The resource provider name 
    85        :return: Return False if we have a reason to believe this didn't work 
    86        :rtype: bool 
    87        """ 
    88        post_url = "{}providers/{}/register?api-version=2016-02-01".format(url_prefix, rp_name) 
    89        get_url = "{}providers/{}?api-version=2016-02-01".format(url_prefix, rp_name) 
    90        _LOGGER.warning( 
    91            "Resource provider '%s' used by this operation is not registered. We are registering for you.", 
    92            rp_name, 
    93        ) 
    94        post_response = await self.next.send(self._build_next_request(initial_request, "POST", post_url)) 
    95        if post_response.http_response.status_code != 200: 
    96            _LOGGER.warning("Registration failed. Please register manually.") 
    97            return False 
    98 
    99        while True: 
    100            await asyncio.sleep(10) 
    101            get_response = await self.next.send(self._build_next_request(initial_request, "GET", get_url)) 
    102            rp_info = json.loads(get_response.http_response.text()) 
    103            if rp_info["registrationState"] == "Registered": 
    104                _LOGGER.warning("Registration succeeded.") 
    105                return True