Skip to content

[Core] Fix get_token() issue in msi login and expiresIn key error in cloud shell login credentials for track 2 SDK related commands #14187

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

Merged
Juliehzl merged 23 commits into from
Jul 9, 2020
Merged

[Core] Fix get_token() issue in msi login and expiresIn key error in cloud shell login credentials for track 2 SDK related commands #14187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
d8a71bf
resolve ketError for Cloud Shell
Juliehzl Jul 2, 2020
7acdefa
fix cloud shell
Jul 2, 2020
e9afbbf
add get_token for msi
Jul 2, 2020
9f2be7a
add more to ignore
Jul 2, 2020
87aa599
reset
Jul 2, 2020
83c1654
ignore launch.json
Jul 2, 2020
141f00d
rename class
Jul 2, 2020
6c41d3c
fix style
Jul 2, 2020
b93e244
pass test_get_login_credentials_msi_system_assigned
Juliehzl Jul 3, 2020
0f54593
pass test_get_login_credentials_msi_system_assigned
Juliehzl Jul 3, 2020
f95b978
❯ azdev test test_get_login_credentials_msi_user_assigned_with_client_id
Juliehzl Jul 3, 2020
d872b65
test_get_login_credentials_msi_user_assigned_with_object_id
Juliehzl Jul 3, 2020
936ab6d
test_get_login_credentials_msi_user_assigned_with_res_id
Juliehzl Jul 3, 2020
e9f5441
test pass get_raaw_token
Juliehzl Jul 3, 2020
e751ece
test_get_raw_token_in_cloud_console
Juliehzl Jul 3, 2020
55b350a
fix test
Juliehzl Jul 3, 2020
41b2ae1
test pass
Juliehzl Jul 3, 2020
4cf1fcb
test pass
Juliehzl Jul 3, 2020
5bdfe74
revert launch.json
Juliehzl Jul 3, 2020
8a5985e
change gitignore
Juliehzl Jul 3, 2020
2ab696b
import as required
Juliehzl Jul 3, 2020
00e01aa
resolve comments
Juliehzl Jul 9, 2020
6085e2a
mock test
Juliehzl Jul 9, 2020
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
23 changes: 11 additions & 12 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 @@ -22,6 +22,8 @@
is_windows, is_wsl
from azure.cli.core.cloud import get_active_cloud, set_cloud_subscription

from .adal_authentication import MSIAuthenticationWrapper

@jiasli jiasli Jul 6, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Another problem with mock.patch is for this form of import, it won't work. After importing MSIAuthenticationWrapper, it becomes a local reference azure.cli.core._profile.MSIAuthenticationWrapper. See Where to patch.


from knack.log import get_logger
from knack.util import CLIError

Expand Down Expand Up @@ -307,18 +309,17 @@ def find_subscriptions_in_vm_with_msi(self, identity_id=None, allow_no_subscript

import jwt
from requests import HTTPError
from msrestazure.azure_active_directory import MSIAuthentication
from msrestazure.tools import is_valid_resource_id
resource = self.cli_ctx.cloud.endpoints.active_directory_resource_id

if identity_id:
if is_valid_resource_id(identity_id):
msi_creds = MSIAuthentication(resource=resource, msi_res_id=identity_id)
msi_creds = MSIAuthenticationWrapper(resource=resource, msi_res_id=identity_id)
identity_type = MsiAccountTypes.user_assigned_resource_id
else:
authenticated = False
try:
msi_creds = MSIAuthentication(resource=resource, client_id=identity_id)
msi_creds = MSIAuthenticationWrapper(resource=resource, client_id=identity_id)
identity_type = MsiAccountTypes.user_assigned_client_id
authenticated = True
except HTTPError as ex:
Expand All @@ -330,7 +331,7 @@ def find_subscriptions_in_vm_with_msi(self, identity_id=None, allow_no_subscript
if not authenticated:
try:
identity_type = MsiAccountTypes.user_assigned_object_id
msi_creds = MSIAuthentication(resource=resource, object_id=identity_id)
msi_creds = MSIAuthenticationWrapper(resource=resource, object_id=identity_id)
authenticated = True
except HTTPError as ex:
if ex.response.reason == 'Bad Request' and ex.response.status == 400:
Expand All @@ -343,7 +344,7 @@ def find_subscriptions_in_vm_with_msi(self, identity_id=None, allow_no_subscript

else:
identity_type = MsiAccountTypes.system_assigned
msi_creds = MSIAuthentication(resource=resource)
msi_creds = MSIAuthenticationWrapper(resource=resource)

token_entry = msi_creds.token
token = token_entry['access_token']
Expand Down Expand Up @@ -388,8 +389,7 @@ def find_subscriptions_in_cloud_console(self):
return deepcopy(consolidated)

def _get_token_from_cloud_shell(self, resource): # pylint: disable=no-self-use
from msrestazure.azure_active_directory import MSIAuthentication
auth = MSIAuthentication(resource=resource)
auth = MSIAuthenticationWrapper(resource=resource)
auth.set_token()
token_entry = auth.token
return (token_entry['token_type'], token_entry['access_token'], token_entry)
Expand Down Expand Up @@ -774,15 +774,14 @@ def valid_msi_account_types():

@staticmethod
def msi_auth_factory(cli_account_name, identity, resource):
from msrestazure.azure_active_directory import MSIAuthentication
if cli_account_name == MsiAccountTypes.system_assigned:
return MSIAuthentication(resource=resource)
return MSIAuthenticationWrapper(resource=resource)
if cli_account_name == MsiAccountTypes.user_assigned_client_id:
return MSIAuthentication(resource=resource, client_id=identity)
return MSIAuthenticationWrapper(resource=resource, client_id=identity)
if cli_account_name == MsiAccountTypes.user_assigned_object_id:
return MSIAuthentication(resource=resource, object_id=identity)
return MSIAuthenticationWrapper(resource=resource, object_id=identity)
if cli_account_name == MsiAccountTypes.user_assigned_resource_id:
return MSIAuthentication(resource=resource, msi_res_id=identity)
return MSIAuthenticationWrapper(resource=resource, msi_res_id=identity)
raise ValueError("unrecognized msi account name '{}'".format(cli_account_name))


Expand Down
14 changes: 12 additions & 2 deletions src/azure-cli-core/azure/cli/core/adal_authentication.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import adal

from msrest.authentication import Authentication
from msrestazure.azure_active_directory import MSIAuthentication
from azure.core.credentials import AccessToken
from azure.cli.core.util import in_cloud_console

Expand Down Expand Up @@ -60,8 +61,10 @@ def _get_token(self):
# This method is exposed for Azure Core.
def get_token(self, *scopes, **kwargs): # pylint:disable=unused-argument
_, token, full_token, _ = self._get_token()

return AccessToken(token, int(full_token['expiresIn'] + time.time()))
try:
return AccessToken(token, int(full_token['expiresIn'] + time.time()))
except KeyError: # needed to deal with differing unserialized MSI token payload
return AccessToken(token, int(full_token['expires_on']))

@jsntcy jsntcy Jul 3, 2020 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This logic is for cloud shell? If yes, can we add comments so that others can understand the purpose?

@Juliehzl Juliehzl Jul 6, 2020

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

sure


# This method is exposed for msrest.
def signed_session(self, session=None): # pylint: disable=arguments-differ
Expand All @@ -83,3 +86,10 @@ def _log_hostname():
logger = get_logger(__name__)
logger.warning("A Cloud Shell credential problem occurred. When you report the issue with the error "
"below, please mention the hostname '%s'", socket.gethostname())


class MSIAuthenticationWrapper(MSIAuthentication):
# This method is exposed for Azure Core.
def get_token(self):
self.set_token()
return AccessToken(self.token['access_token'], int(self.token['expires_on']))
16 changes: 8 additions & 8 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 @@ -645,7 +645,7 @@ def test_get_login_credentials_aux_tenants(self, mock_get_token, mock_read_cred_
aux_tenants=[test_tenant_id2])

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)

@jiasli jiasli Jul 3, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Is this change necessary? Won't it be nice to have MSIAuthenticationWrapper also tested?

@jiasli jiasli Jul 6, 2020 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

mock.patch by default will create a MagicMock instance which can't be inherited, maybe using new=MSIAuthentication may work? Haven't tested.

def test_get_login_credentials_msi_system_assigned(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -676,7 +676,7 @@ def test_get_login_credentials_msi_system_assigned(self, mock_msi_auth, mock_rea
self.assertTrue(cred.token_read_count)

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_login_credentials_msi_user_assigned_with_client_id(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -707,7 +707,7 @@ def test_get_login_credentials_msi_user_assigned_with_client_id(self, mock_msi_a
self.assertTrue(cred.client_id, test_client_id)

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_login_credentials_msi_user_assigned_with_object_id(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -738,7 +738,7 @@ def test_get_login_credentials_msi_user_assigned_with_object_id(self, mock_msi_a
self.assertTrue(cred.object_id, test_object_id)

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_login_credentials_msi_user_assigned_with_res_id(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -849,7 +849,7 @@ def test_get_raw_token_for_sp(self, mock_get_token, mock_read_cred_file):
self.assertEqual(tenant, self.tenant_id)

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_raw_token_msi_system_assigned(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -884,7 +884,7 @@ def test_get_raw_token_msi_system_assigned(self, mock_msi_auth, mock_read_cred_f

@mock.patch('azure.cli.core._profile.in_cloud_console', autospec=True)
@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_raw_token_in_cloud_console(self, mock_msi_auth, mock_read_cred_file, mock_in_cloud_console):
mock_read_cred_file.return_value = []
mock_in_cloud_console.return_value = True
Expand Down Expand Up @@ -1037,7 +1037,7 @@ def test_find_subscriptions_thru_username_non_password(self, mock_auth_context):
# assert
self.assertEqual([], subs)

@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
@mock.patch('azure.cli.core.profiles._shared.get_client_class', autospec=True)
@mock.patch('azure.cli.core._profile._get_cloud_console_token_endpoint', autospec=True)
@mock.patch('azure.cli.core._profile.SubscriptionFinder', autospec=True)
Expand Down Expand Up @@ -1186,7 +1186,7 @@ def __init__(self, *args, **kwargs):
self.assertEqual(s['id'], self.id1.split('/')[-1])
self.assertEqual(s['tenantId'], '54826b22-38d6-4fb2-bad9-b7b93a3e9c5a')

@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
@mock.patch('azure.cli.core.profiles._shared.get_client_class', autospec=True)
@mock.patch('azure.cli.core._profile.SubscriptionFinder', autospec=True)
def test_find_subscriptions_in_vm_with_msi_user_assigned_with_object_id(self, mock_subscription_finder, mock_get_client_class,
Expand Down
14 changes: 7 additions & 7 deletions src/azure-cli-core/azure/cli/core/tests/test_profile_v2016_06_01.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -571,7 +571,7 @@ def test_get_login_credentials_aux_subscriptions(self, mock_get_token, mock_read
self.assertEqual(mock_get_token.call_count, 2)

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_login_credentials_msi_system_assigned(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -602,7 +602,7 @@ def test_get_login_credentials_msi_system_assigned(self, mock_msi_auth, mock_rea
self.assertTrue(cred.token_read_count)

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_login_credentials_msi_user_assigned_with_client_id(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -633,7 +633,7 @@ def test_get_login_credentials_msi_user_assigned_with_client_id(self, mock_msi_a
self.assertTrue(cred.client_id, test_client_id)

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_login_credentials_msi_user_assigned_with_object_id(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -664,7 +664,7 @@ def test_get_login_credentials_msi_user_assigned_with_object_id(self, mock_msi_a
self.assertTrue(cred.object_id, test_object_id)

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_login_credentials_msi_user_assigned_with_res_id(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -753,7 +753,7 @@ def test_get_raw_token_for_sp(self, mock_get_token, mock_read_cred_file):
self.assertEqual(tenant, self.tenant_id)

@mock.patch('azure.cli.core._profile._load_tokens_from_file', autospec=True)
@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
def test_get_raw_token_msi_system_assigned(self, mock_msi_auth, mock_read_cred_file):
mock_read_cred_file.return_value = []

Expand Down Expand Up @@ -899,7 +899,7 @@ def test_find_subscriptions_thru_username_non_password(self, mock_auth_context):
# assert
self.assertEqual([], subs)

@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
@mock.patch('azure.cli.core.profiles._shared.get_client_class', autospec=True)
@mock.patch('azure.cli.core._profile._get_cloud_console_token_endpoint', autospec=True)
@mock.patch('azure.cli.core._profile.SubscriptionFinder', autospec=True)
Expand Down Expand Up @@ -1048,7 +1048,7 @@ def __init__(self, *args, **kwargs):
self.assertEqual(s['id'], self.id1.split('/')[-1])
self.assertEqual(s['tenantId'], '54826b22-38d6-4fb2-bad9-b7b93a3e9c5a')

@mock.patch('msrestazure.azure_active_directory.MSIAuthentication', autospec=True)
@mock.patch('azure.cli.core._profile.MSIAuthenticationWrapper', autospec=True)
@mock.patch('azure.cli.core.profiles._shared.get_client_class', autospec=True)
@mock.patch('azure.cli.core._profile.SubscriptionFinder', autospec=True)
def test_find_subscriptions_in_vm_with_msi_user_assigned_with_object_id(self, mock_subscription_finder, mock_get_client_class,
Expand Down