Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

{Profile} az account get-access-token: Show expiresOn for managed identity #20219

Merged
merged 10 commits into from
Nov 17, 2021
Merged

{Profile} az account get-access-token: Show expiresOn for managed identity #20219

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
bfaacc2
expires_on
jiasli Nov 5, 2021
5859d97
unittest
jiasli Nov 5, 2021
71d161e
datetime
jiasli Nov 8, 2021
141ab62
Merge remote-tracking branch 'azure/dev' into expires-on
jiasli Nov 8, 2021
f65d76a
resouce
jiasli Nov 8, 2021
4cf877e
Add tests
jiasli Nov 8, 2021
173d469
comment
jiasli Nov 15, 2021
7aa714e
Merge remote-tracking branch 'azure/dev' into expires-on
jiasli Nov 15, 2021
eae0254
Merge remote-tracking branch 'azure/dev' into expires-on
jiasli Nov 16, 2021
f0746f0
refine tests
jiasli Nov 16, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 22 additions & 25 deletions src/azure-cli-core/azure/cli/core/_profile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -361,37 +361,41 @@ def get_raw_token(self, resource=None, scopes=None, subscription=None, tenant=No
raise CLIError("Please specify only one of subscription and tenant, not both")

account = self.get_subscription(subscription)
resource = resource or self.cli_ctx.cloud.endpoints.active_directory_resource_id

# This is not used anyway, just a placeholder
mi_resoure = self.cli_ctx.cloud.endpoints.active_directory_resource_id

identity_type, identity_id = Profile._try_parse_msi_account_name(account)
if identity_type:
# MSI
# managed identity
if tenant:
raise CLIError("Tenant shouldn't be specified for MSI account")
msi_creds = MsiAccountTypes.msi_auth_factory(identity_type, identity_id, resource)
msi_creds.set_token()
token_entry = msi_creds.token
creds = (token_entry['token_type'], token_entry['access_token'], token_entry)
raise CLIError("Tenant shouldn't be specified for managed identity account")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we use a specific error type (such as ArgumentUsageError) instead of CLIError?

Copy link
Member Author

@jiasli jiasli Nov 17, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are so many places in core that don't comply with the error handling rule. Let's refine them together later.

msi_creds = MsiAccountTypes.msi_auth_factory(identity_type, identity_id, mi_resoure)
sdk_token = msi_creds.get_token(*scopes)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The old implementation of ADAL-based Azure CLI returns token_entry containing inconsistent fields:

  • expires_on for managed identity
  • expiresOn for ADAL credential

They are unified later at command module level:

azure-cli/src/azure-cli/azure/cli/command_modules/profile/custom.py

Lines 82 to 85 in df737ed

if 'expires_on' in token_entry:
# https://docs.python.org/3.8/library/datetime.html#strftime-and-strptime-format-codes
token_entry['expiresOn'] = _fromtimestamp(int(token_entry['expires_on']))\
.strftime("%Y-%m-%d %H:%M:%S.%f")

We use get_token to unify them to epoch int expires_on in core instead.

Copy link
Contributor

@zhoxing-ms zhoxing-ms Nov 17, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just out of curiosity, will token_entry['expiresOn'] has its own value in some cases before it is overwritten by _fromtimestamp(int(token_entry['expires_on'])).strftime("%Y-%m-%d %H:%M:%S.%f")?
May I ask their values should be the same in all cases, but only in different formats, right?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No. get_token only returns AccessToken which must have int epoch expires_on.

elif in_cloud_console() and account[_USER_ENTITY].get(_CLOUD_SHELL_ID):
# Cloud Shell
if tenant:
raise CLIError("Tenant shouldn't be specified for Cloud Shell account")
creds = self._get_token_from_cloud_shell(resource)
msi_creds = MsiAccountTypes.msi_auth_factory(MsiAccountTypes.system_assigned, identity_id, mi_resoure)
sdk_token = msi_creds.get_token(*scopes)
else:
credential = self._create_credential(account, tenant)
token = credential.get_token(*scopes)
sdk_token = credential.get_token(*scopes)

import datetime
expiresOn = datetime.datetime.fromtimestamp(token.expires_on).strftime("%Y-%m-%d %H:%M:%S.%f")
# Convert epoch int 'expires_on' to datetime string 'expiresOn' for backward compatibility
# WARNING: expiresOn is deprecated and will be removed in future release.
import datetime
expiresOn = datetime.datetime.fromtimestamp(sdk_token.expires_on).strftime("%Y-%m-%d %H:%M:%S.%f")

token_entry = {
'accessToken': token.token,
'expires_on': token.expires_on,
'expiresOn': expiresOn
}
token_entry = {
'accessToken': sdk_token.token,
'expires_on': sdk_token.expires_on, # epoch int, like 1605238724
'expiresOn': expiresOn # datetime string, like "2020-11-12 13:50:47.114324"
}

# (tokenType, accessToken, tokenEntry)
creds = 'Bearer', sdk_token.token, token_entry

# (tokenType, accessToken, tokenEntry)
creds = 'Bearer', token.token, token_entry
# (cred, subscription, tenant)
return (creds,
None if tenant else str(account[_SUBSCRIPTION_ID]),
Expand Down Expand Up @@ -695,13 +699,6 @@ def get_installation_id(self):
self._storage[_INSTALLATION_ID] = installation_id
return installation_id

def _get_token_from_cloud_shell(self, resource): # pylint: disable=no-self-use
from azure.cli.core.auth.adal_authentication import MSIAuthenticationWrapper
auth = MSIAuthenticationWrapper(resource=resource)
auth.set_token()
token_entry = auth.token
return (token_entry['token_type'], token_entry['access_token'], token_entry)


class MsiAccountTypes:
# pylint: disable=no-method-argument,no-self-argument
Expand Down
3 changes: 1 addition & 2 deletions src/azure-cli-core/azure/cli/core/auth/adal_authentication.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,10 @@
# --------------------------------------------------------------------------------------------

import requests
from azure.core.credentials import AccessToken
from knack.log import get_logger
from msrestazure.azure_active_directory import MSIAuthentication

from .util import _normalize_scopes, scopes_to_resource
from .util import _normalize_scopes, scopes_to_resource, AccessToken

logger = get_logger(__name__)

Expand Down
11 changes: 8 additions & 3 deletions src/azure-cli-core/azure/cli/core/auth/msal_authentication.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from knack.util import CLIError
from msal import PublicClientApplication, ConfidentialClientApplication

from .util import check_result
from .util import check_result, AccessToken

# OAuth 2.0 client credentials flow parameter
# https://docs.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-client-creds-grant-flow
Expand Down Expand Up @@ -134,9 +134,14 @@ def _build_sdk_access_token(token_entry):
import time
request_time = int(time.time())

# MSAL token entry sample:
# {
# 'access_token': 'eyJ0eXAiOiJKV...',
# 'token_type': 'Bearer',
# 'expires_in': 1618
# }

# Importing azure.core.credentials.AccessToken is expensive.
# This can slow down commands that doesn't need azure.core, like `az account get-access-token`.
# So We define our own AccessToken.
from collections import namedtuple
AccessToken = namedtuple("AccessToken", ["token", "expires_on"])
return AccessToken(token_entry["access_token"], request_time + token_entry["expires_in"])
2 changes: 1 addition & 1 deletion src/azure-cli-core/azure/cli/core/auth/tests/test_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
# pylint: disable=protected-access

import unittest
from ..util import scopes_to_resource, resource_to_scopes, _normalize_scopes, _generate_login_command
from azure.cli.core.auth.util import scopes_to_resource, resource_to_scopes, _normalize_scopes, _generate_login_command


class TestUtil(unittest.TestCase):
Expand Down
5 changes: 5 additions & 0 deletions src/azure-cli-core/azure/cli/core/auth/util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,16 @@
# --------------------------------------------------------------------------------------------

import os
from collections import namedtuple

from knack.log import get_logger

logger = get_logger(__name__)


AccessToken = namedtuple("AccessToken", ["token", "expires_on"])


def aad_error_handler(error, **kwargs):
""" Handle the error from AAD server returned by ADAL or MSAL. """

Expand Down
51 changes: 30 additions & 21 deletions src/azure-cli-core/azure/cli/core/tests/test_profile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,31 +5,24 @@

# pylint: disable=protected-access
import json
import os
import sys
import unittest
from unittest import mock
import re
import datetime

from copy import deepcopy

from azure.core.credentials import AccessToken
from unittest import mock

from azure.cli.core._profile import (Profile, SubscriptionFinder, _attach_token_tenant,
_transform_subscription_for_multiapi)

from azure.mgmt.resource.subscriptions.models import \
(Subscription, SubscriptionPolicies, SpendingLimit, ManagedByTenant, TenantIdDescription)

from azure.cli.core.auth.util import AccessToken
from azure.cli.core.mock import DummyCli
from azure.identity import AuthenticationRecord
from azure.mgmt.resource.subscriptions.models import \
(Subscription, SubscriptionPolicies, SpendingLimit, ManagedByTenant)

from knack.util import CLIError


MOCK_ACCESS_TOKEN = "mock_access_token"
MOCK_EXPIRES_ON = 1630920323
MOCK_EXPIRES_ON_STR = "1630920323"
MOCK_EXPIRES_ON_INT = 1630920323
MOCK_EXPIRES_ON_DATETIME = '2021-09-06 17:25:23.000000'
BEARER = 'Bearer'


Expand All @@ -43,14 +36,15 @@ def get_token(self, *scopes, **kwargs):
import time
now = int(time.time())
# Mock sdk/identity/azure-identity/azure/identity/_internal/msal_credentials.py:230
return AccessToken(MOCK_ACCESS_TOKEN, MOCK_EXPIRES_ON)
return AccessToken(MOCK_ACCESS_TOKEN, MOCK_EXPIRES_ON_INT)


class MSRestAzureAuthStub:
def __init__(self, *args, **kwargs):
self._token = {
'token_type': 'Bearer',
'access_token': TestProfile.test_msi_access_token
'access_token': TestProfile.test_msi_access_token,
'expires_on': MOCK_EXPIRES_ON_STR
}
self.set_token_invoked_count = 0
self.token_read_count = 0
Expand All @@ -70,6 +64,9 @@ def token(self):
def token(self, value):
self._token = value

def get_token(self, *args, **kwargs):
return AccessToken(self.token['access_token'], int(self.token['expires_on']))


class TestProfile(unittest.TestCase):

Expand Down Expand Up @@ -1049,7 +1046,8 @@ def test_get_raw_token(self):

self.assertEqual(creds[0], 'Bearer')
self.assertEqual(creds[1], MOCK_ACCESS_TOKEN)
self.assertEqual(creds[2]['expires_on'], MOCK_EXPIRES_ON)
self.assertEqual(creds[2]['expires_on'], MOCK_EXPIRES_ON_INT)
self.assertEqual(creds[2]['expiresOn'], MOCK_EXPIRES_ON_DATETIME)

# subscription should be set
self.assertEqual(sub, self.subscription1.subscription_id)
Expand All @@ -1060,7 +1058,8 @@ def test_get_raw_token(self):

self.assertEqual(creds[0], 'Bearer')
self.assertEqual(creds[1], MOCK_ACCESS_TOKEN)
self.assertEqual(creds[2]['expires_on'], MOCK_EXPIRES_ON)
self.assertEqual(creds[2]['expires_on'], MOCK_EXPIRES_ON_INT)
self.assertEqual(creds[2]['expiresOn'], MOCK_EXPIRES_ON_DATETIME)

# subscription shouldn't be set
self.assertIsNone(sub)
Expand All @@ -1084,7 +1083,8 @@ def test_get_raw_token_for_sp(self, get_service_principal_credential_mock):
self.assertEqual(creds[0], BEARER)
self.assertEqual(creds[1], MOCK_ACCESS_TOKEN)
# the last in the tuple is the whole token entry which has several fields
self.assertEqual(creds[2]['expires_on'], MOCK_EXPIRES_ON)
self.assertEqual(creds[2]['expires_on'], MOCK_EXPIRES_ON_INT)
self.assertEqual(creds[2]['expiresOn'], MOCK_EXPIRES_ON_DATETIME)

# subscription should be set
self.assertEqual(sub, self.subscription1.subscription_id)
Expand All @@ -1095,7 +1095,8 @@ def test_get_raw_token_for_sp(self, get_service_principal_credential_mock):

self.assertEqual(creds[0], BEARER)
self.assertEqual(creds[1], MOCK_ACCESS_TOKEN)
self.assertEqual(creds[2]['expires_on'], MOCK_EXPIRES_ON)
self.assertEqual(creds[2]['expires_on'], MOCK_EXPIRES_ON_INT)
self.assertEqual(creds[2]['expiresOn'], MOCK_EXPIRES_ON_DATETIME)

# subscription shouldn't be set
self.assertIsNone(sub)
Expand Down Expand Up @@ -1124,11 +1125,15 @@ def test_get_raw_token_msi_system_assigned(self, mock_msi_auth):
self.assertEqual(subscription_id, test_subscription_id)
self.assertEqual(cred[0], 'Bearer')
self.assertEqual(cred[1], TestProfile.test_msi_access_token)

# Make sure expires_on and expiresOn are set
self.assertEqual(cred[2]['expires_on'], MOCK_EXPIRES_ON_INT)
self.assertEqual(cred[2]['expiresOn'], MOCK_EXPIRES_ON_DATETIME)
self.assertEqual(subscription_id, test_subscription_id)
self.assertEqual(tenant_id, test_tenant_id)

# verify tenant shouldn't be specified for MSI account
with self.assertRaisesRegexp(CLIError, "MSI"):
with self.assertRaisesRegexp(CLIError, "Tenant shouldn't be specified"):
cred, subscription_id, _ = profile.get_raw_token(resource='http://test_resource', tenant=self.tenant_id)

@mock.patch('azure.cli.core._profile.in_cloud_console', autospec=True)
Expand Down Expand Up @@ -1157,6 +1162,10 @@ def test_get_raw_token_in_cloud_console(self, mock_msi_auth, mock_in_cloud_conso
self.assertEqual(subscription_id, test_subscription_id)
self.assertEqual(cred[0], 'Bearer')
self.assertEqual(cred[1], TestProfile.test_msi_access_token)

# Make sure expires_on and expiresOn are set
self.assertEqual(cred[2]['expires_on'], MOCK_EXPIRES_ON_INT)
self.assertEqual(cred[2]['expiresOn'], MOCK_EXPIRES_ON_DATETIME)
self.assertEqual(subscription_id, test_subscription_id)
self.assertEqual(tenant_id, test_tenant_id)

Expand Down
2 changes: 1 addition & 1 deletion src/azure-cli/azure/cli/command_modules/profile/custom.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ def get_access_token(cmd, subscription=None, resource=None, scopes=None, resourc
'tokenType': creds[0],
'accessToken': creds[1],
# 'expires_on': creds[2].get('expires_on', None),
'expiresOn': creds[2].get('expiresOn', None),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do you need to change this line?

Copy link
Member Author

@jiasli jiasli Nov 8, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need to guarantee that expiresOn always exists in creds[2] - CLI should fail if expiresOn is not set, instead of returning None which will cause more trouble.

'expiresOn': creds[2]['expiresOn'],
'tenant': tenant
}
if subscription:
Expand Down