Skip to content

Instantly share code, notes, and snippets.

@lmazuel
Last active December 6, 2023 14:47
Show Gist options
  • Save lmazuel/cc683d82ea1d7b40208de7c9fc8de59d to your computer and use it in GitHub Desktop.
Save lmazuel/cc683d82ea1d7b40208de7c9fc8de59d to your computer and use it in GitHub Desktop.
Use azure-identity to authenticate mgmt SDKs that needs "msrestazure/azure.common.credentials" protocol
# Wrap credentials from azure-identity to be compatible with SDK that needs msrestazure or azure.common.credentials
# Need msrest >= 0.6.0
# See also https://pypi.org/project/azure-identity/
from msrest.authentication import BasicTokenAuthentication
from azure.core.pipeline.policies import BearerTokenCredentialPolicy
from azure.core.pipeline import PipelineRequest, PipelineContext
from azure.core.pipeline.transport import HttpRequest
from azure.identity import DefaultAzureCredential
class CredentialWrapper(BasicTokenAuthentication):
def __init__(self, credential=None, resource_id="https://management.azure.com/.default", **kwargs):
"""Wrap any azure-identity credential to work with SDK that needs azure.common.credentials/msrestazure.
Default resource is ARM (syntax of endpoint v2)
:param credential: Any azure-identity credential (DefaultAzureCredential by default)
:param str resource_id: The scope to use to get the token (default ARM)
"""
super(CredentialWrapper, self).__init__(None)
if credential is None:
credential = DefaultAzureCredential()
self._policy = BearerTokenCredentialPolicy(credential, resource_id, **kwargs)
def _make_request(self):
return PipelineRequest(
HttpRequest(
"CredentialWrapper",
"https://fakeurl"
),
PipelineContext(None)
)
def set_token(self):
"""Ask the azure-core BearerTokenCredentialPolicy policy to get a token.
Using the policy gives us for free the caching system of azure-core.
We could make this code simpler by using private method, but by definition
I can't assure they will be there forever, so mocking a fake call to the policy
to extract the token, using 100% public API."""
request = self._make_request()
self._policy.on_request(request)
# Read Authorization, and get the second part after Bearer
token = request.http_request.headers["Authorization"].split(" ", 1)[1]
self.token = {"access_token": token}
def signed_session(self, session=None):
self.set_token()
return super(CredentialWrapper, self).signed_session(session)
if __name__ == "__main__":
import os
credentials = CredentialWrapper()
subscription_id = os.environ.get("AZURE_SUBSCRIPTION_ID", "<subscription_id>")
from azure.mgmt.resource import ResourceManagementClient
client = ResourceManagementClient(credentials, subscription_id)
for rg in client.resource_groups.list():
print(rg.name)
@jakevc
Copy link

jakevc commented Mar 9, 2023

@lmazuel
Copy link
Author

lmazuel commented Mar 10, 2023

Yes, that's correct, Batch is in the process of doing it, they are actively working on it. I forgot batch I admit :)

@jakevc
Copy link

jakevc commented Mar 10, 2023

Cool thanks for confirmation, can you link to the PR? I'm using your workaround, but would love to get rid of the msrest dependency in
from msrest.authentication import BasicTokenAuthentication

@vitalybondarenko
Copy link

Hello! i still need to use CredentialWrapper with
azure.mgmt.sqlvirtualmachine.SqlVirtualMachineManagementClient and
azure.mgmt.documentdb.DocumentDB
Are they planned for update?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment