codekingpro/portable-devtools

1# ------------------------------------2# Copyright (c) Microsoft Corporation.3# Licensed under the MIT License.4# ------------------------------------5import time6from typing import Iterable, Union, Optional, Any7 8from azure.core.credentials import AccessToken9from azure.core.pipeline import Pipeline10from azure.core.pipeline.transport import HttpRequest11from .aad_client_base import AadClientBase12from .aadclient_certificate import AadClientCertificate13from .pipeline import build_pipeline14 15 16class AadClient(AadClientBase):17    def __enter__(self):18        self._pipeline.__enter__()19        return self20 21    def __exit__(self, *args):22        self._pipeline.__exit__(*args)23 24    def close(self) -> None:25        self.__exit__()26 27    def obtain_token_by_authorization_code(28        self, scopes: Iterable[str], code: str, redirect_uri: str, client_secret: Optional[str] = None, **kwargs: Any29    ) -> AccessToken:30        request = self._get_auth_code_request(31            scopes=scopes, code=code, redirect_uri=redirect_uri, client_secret=client_secret, **kwargs32        )33        return self._run_pipeline(request, **kwargs)34 35    def obtain_token_by_client_certificate(36        self, scopes: Iterable[str], certificate: AadClientCertificate, **kwargs: Any37    ) -> AccessToken:38        request = self._get_client_certificate_request(scopes, certificate, **kwargs)39        return self._run_pipeline(request, **kwargs)40 41    def obtain_token_by_client_secret(self, scopes: Iterable[str], secret: str, **kwargs: Any) -> AccessToken:42        request = self._get_client_secret_request(scopes, secret, **kwargs)43        return self._run_pipeline(request, **kwargs)44 45    def obtain_token_by_jwt_assertion(self, scopes: Iterable[str], assertion: str, **kwargs: Any) -> AccessToken:46        request = self._get_jwt_assertion_request(scopes, assertion, **kwargs)47        return self._run_pipeline(request, **kwargs)48 49    def obtain_token_by_refresh_token(self, scopes: Iterable[str], refresh_token: str, **kwargs: Any) -> AccessToken:50        request = self._get_refresh_token_request(scopes, refresh_token, **kwargs)51        return self._run_pipeline(request, **kwargs)52 53    def obtain_token_on_behalf_of(54        self,55        scopes: Iterable[str],56        client_credential: Union[str, AadClientCertificate],57        user_assertion: str,58        **kwargs: Any59    ) -> AccessToken:60        # no need for an implementation, non-async OnBehalfOfCredential acquires tokens through MSAL61        raise NotImplementedError()62 63    def _build_pipeline(self, **kwargs: Any) -> Pipeline:64        return build_pipeline(**kwargs)65 66    def _run_pipeline(self, request: HttpRequest, **kwargs: Any) -> AccessToken:67        # remove tenant_id and claims kwarg that could have been passed from credential's get_token method68        # tenant_id is already part of `request` at this point69        kwargs.pop("tenant_id", None)70        kwargs.pop("claims", None)71        enable_cae = kwargs.pop("enable_cae", False)72        now = int(time.time())73        response = self._pipeline.run(request, retry_on_methods=self._POST, **kwargs)74        return self._process_response(response, now, enable_cae=enable_cae, **kwargs)75