def get_aad_token()

in cosmos/scripts/chaos/cosmosdb_client_network_chaos_tool/get_cdb_aad_token.py [0:0]

• 23 lines of code
• 10 McCabe index (conditional complexity)

def get_aad_token(endpoint, client_id, client_secret, tenant_id):
    """:
    This function returns a Microsoft Entra ID token for the given endpoint using the given 
        1. client_id, client_secret and tenant_id combination when using a service principal
        2. client_id only when using a managed identity

    Args:
        endpoint (str): The endpoint for which to retrieve the Microsoft Entra ID token.
        client_id (str): The client ID of the Managed Identity or the Microsoft Entra ID application.
        client_secret (str): The client secret of the Microsoft Entra ID application.
        tenant_id (str): The ID of the Microsoft Entra ID tenant.

    Returns:
        str: The Microsoft Entra ID token.

    Raises:
        Exception: If the Microsoft Entra ID token cannot be retrieved.
    """
    max_retries = 3
    retry_delay = 5

    for retry in range(max_retries):
        try:
            if client_id and client_secret and tenant_id:
                aad_credentials = ClientSecretCredential(tenant_id, client_id, client_secret)
            elif client_id and not client_secret and not tenant_id:
                aad_credentials = DefaultAzureCredential(managed_identity_client_id=client_id)
            else:
                raise Exception("Either provide Client ID only to retrieve the Microsoft Entra ID token using Manged Identity or provide Client ID, Client Secret and Tenant ID to retrieve the Microsoft Entra ID token using Service Principal.")
            
            result = endpoint.split(':')
            scope = result[0] + ":" + result[1] + "/.default"
            token = aad_credentials.get_token(scope)
            print(token.token)
            break
        except Exception as e:
            print("Error occurred while retrieving the Microsoft Entra ID token:", str(e))
            if retry < max_retries - 1:
                print(f"Retrying in {retry_delay} seconds...")
                time.sleep(retry_delay)
            else:
                raise Exception("Failed to retrieve Microsoft Entra ID token after multiple retries.")