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.set_token()
token_entry = auth.token
return (token_entry['token_type'], token_entry['access_token'],
token_entry)
def test_msi_vm_imds_no_retry_on_bad_error(self):
httpretty.register_uri(httpretty.GET,
'http://169.254.169.254/metadata/identity/oauth2/token',
status=499)
credentials = MSIAuthentication()
with self.assertRaises(HTTPError) as cm:
credentials.set_token()
def find_subscriptions_in_vm_with_msi(self, identity_id=None):
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
msi_creds = MSIAuthentication()
token_entry = None
if identity_id:
if is_valid_resource_id(identity_id):
msi_creds = MSIAuthentication(resource=resource, msi_res_id=identity_id)
identity_type = MsiAccountTypes.user_assigned_resource_id
else:
msi_creds = MSIAuthentication(resource=resource, client_id=identity_id)
try:
msi_creds.set_token()
token_entry = msi_creds.token
identity_type = MsiAccountTypes.user_assigned_client_id
except HTTPError as ex:
if ex.response.reason == 'Bad Request' and ex.response.status == 400:
identity_type = MsiAccountTypes.user_assigned_object_id
msi_creds = MSIAuthentication(resource=resource, object_id=identity_id)
else:
raise
else:
identity_type = MsiAccountTypes.system_assigned
msi_creds = MSIAuthentication(resource=resource)
if not token_entry:
msi_creds.set_token()
token_entry = msi_creds.token
token = token_entry['access_token']
logger.info('MSI: token was retrieved. Now trying to initialize local accounts...')
decode = jwt.decode(token, verify=False, algorithms=['RS256'])
tenant = decode['tid']
subscription_finder = SubscriptionFinder(self.cli_ctx, self.auth_ctx_factory, None)
subscriptions = subscription_finder.find_from_raw_token(tenant, token)
if not subscriptions:
raise CLIError('No access was configured for the VM, hence no subscriptions were found')
base_name = ('{}-{}'.format(identity_type, identity_id) if identity_id else identity_type)
user = _USER_ASSIGNED_IDENTITY if identity_id else _SYSTEM_ASSIGNED_IDENTITY
consolidated = self._normalize_properties(user, subscriptions, is_service_principal=True)
for s in consolidated:
s[_SUBSCRIPTION_NAME] = base_name
# key-off subscription name to allow accounts with same id(but under different identities)
self._set_subscriptions(consolidated, secondary_key_name=_SUBSCRIPTION_NAME)
return deepcopy(consolidated)
def find_subscriptions_in_vm_with_msi(self, identity_id=None):
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
msi_creds = MSIAuthentication()
token_entry = None
if identity_id:
if is_valid_resource_id(identity_id):
msi_creds = MSIAuthentication(resource=resource, msi_res_id=identity_id)
identity_type = MsiAccountTypes.user_assigned_resource_id
else:
msi_creds = MSIAuthentication(resource=resource, client_id=identity_id)
try:
msi_creds.set_token()
token_entry = msi_creds.token
identity_type = MsiAccountTypes.user_assigned_client_id
except HTTPError as ex:
if ex.response.reason == 'Bad Request' and ex.response.status == 400:
identity_type = MsiAccountTypes.user_assigned_object_id
msi_creds = MSIAuthentication(resource=resource, object_id=identity_id)
else:
raise
else:
identity_type = MsiAccountTypes.system_assigned
msi_creds = MSIAuthentication(resource=resource)
if not token_entry:
msi_creds.set_token()
token_entry = msi_creds.token
token = token_entry['access_token']
logger.info('MSI: token was retrieved. Now trying to initialize local accounts...')
decode = jwt.decode(token, verify=False, algorithms=['RS256'])
tenant = decode['tid']
subscription_finder = SubscriptionFinder(self.cli_ctx, self.auth_ctx_factory, None)
subscriptions = subscription_finder.find_from_raw_token(tenant, token)
if not subscriptions:
raise CLIError('No access was configured for the VM, hence no subscriptions were found')
base_name = ('{}-{}'.format(identity_type, identity_id) if identity_id else identity_type)
user = _USER_ASSIGNED_IDENTITY if identity_id else _SYSTEM_ASSIGNED_IDENTITY
consolidated = self._normalize_properties(user, subscriptions, is_service_principal=True)
for s in consolidated:
s[_SUBSCRIPTION_NAME] = base_name
# key-off subscription name to allow accounts with same id(but under different identities)
self._set_subscriptions(consolidated, secondary_key_name=_SUBSCRIPTION_NAME)
return deepcopy(consolidated)
def test_msi_vm_imds_retry(self):
json_payload = {
'token_type': "TokenTypeIMDS",
"access_token": "AccessToken"
}
httpretty.register_uri(httpretty.GET,
'http://169.254.169.254/metadata/identity/oauth2/token',
status=404)
httpretty.register_uri(httpretty.GET,
'http://169.254.169.254/metadata/identity/oauth2/token',
status=429)
httpretty.register_uri(httpretty.GET,
'http://169.254.169.254/metadata/identity/oauth2/token',
status=599)
httpretty.register_uri(httpretty.GET,
'http://169.254.169.254/metadata/identity/oauth2/token',
body=json.dumps(json_payload),
content_type="application/json")
credentials = MSIAuthentication()
credentials.set_token()
assert credentials.scheme == "TokenTypeIMDS"
assert credentials.token == json_payload
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.set_token()
token_entry = auth.token
return (token_entry['token_type'], token_entry['access_token'], token_entry)
class _AadHelper:
authentication_method = None
auth_context = None
msi_auth_context = None
username = None
kusto_uri = None
authority_uri = None
client_id = None
password = None
thumbprint = None
certificate = None
msi_params = None
token_provider = None
def __init__(self, kcsb):
self.kusto_uri = "{0.scheme}://{0.hostname}".format(urlparse(kcsb.data_source))
self.username = None
if all([kcsb.aad_user_id, kcsb.password]):
self.authentication_method = AuthenticationMethod.aad_username_password
self.client_id = "db662dc1-0cfe-4e1c-a843-19a68e65be58"
self.username = kcsb.aad_user_id
self.password = kcsb.password
elif all([kcsb.application_client_id, kcsb.application_key]):
self.authentication_method = AuthenticationMethod.aad_application_key
self.client_id = kcsb.application_client_id
self.client_secret = kcsb.application_key
elif all([kcsb.application_client_id, kcsb.application_certificate, kcsb.application_certificate_thumbprint]):
self.authentication_method = AuthenticationMethod.aad_application_certificate
self.client_id = kcsb.application_client_id
self.certificate = kcsb.application_certificate
self.thumbprint = kcsb.application_certificate_thumbprint
elif kcsb.msi_authentication:
self.authentication_method = AuthenticationMethod.aad_msi
self.msi_params = kcsb.msi_parameters
return
elif any([kcsb.user_token, kcsb.application_token]):
self.token = kcsb.user_token or kcsb.application_token
self.authentication_method = AuthenticationMethod.aad_token
return
elif kcsb.az_cli:
self.authentication_method = AuthenticationMethod.az_cli_profile
return
elif kcsb.token_provider:
self.authentication_method = AuthenticationMethod.token_provider
self.token_provider = kcsb.token_provider
else:
self.authentication_method = AuthenticationMethod.aad_device_login
self.client_id = "db662dc1-0cfe-4e1c-a843-19a68e65be58"
authority = kcsb.authority_id or "common"
aad_authority_uri = os.environ.get("AadAuthorityUri", CLOUD_LOGIN_URL)
self.authority_uri = aad_authority_uri + authority if aad_authority_uri.endswith("/") else aad_authority_uri + "/" + authority
def acquire_authorization_header(self):
"""Acquire tokens from AAD."""
try:
return self._acquire_authorization_header()
except (AdalError, KustoClientError) as error:
if self.authentication_method is AuthenticationMethod.aad_username_password:
kwargs = {"username": self.username, "client_id": self.client_id}
elif self.authentication_method is AuthenticationMethod.aad_application_key:
kwargs = {"client_id": self.client_id}
elif self.authentication_method is AuthenticationMethod.aad_device_login:
kwargs = {"client_id": self.client_id}
elif self.authentication_method is AuthenticationMethod.aad_application_certificate:
kwargs = {"client_id": self.client_id, "thumbprint": self.thumbprint}
elif self.authentication_method is AuthenticationMethod.aad_msi:
kwargs = self.msi_params
elif self.authentication_method is AuthenticationMethod.token_provider:
kwargs = {}
else:
raise error
kwargs["resource"] = self.kusto_uri
if self.authentication_method is AuthenticationMethod.aad_msi:
kwargs["authority"] = AuthenticationMethod.aad_msi.value
elif self.authentication_method is AuthenticationMethod.token_provider:
kwargs["authority"] = AuthenticationMethod.token_provider.value
elif self.auth_context is not None:
kwargs["authority"] = self.auth_context.authority.url
raise KustoAuthenticationError(self.authentication_method.value, error, **kwargs)
def _acquire_authorization_header(self) -> str:
# Token was provided by caller
if self.authentication_method is AuthenticationMethod.aad_token:
return _get_header("Bearer", self.token)
if self.authentication_method is AuthenticationMethod.token_provider:
caller_token = self.token_provider()
if not isinstance(caller_token, str):
raise KustoClientError("Token provider returned something that is not a string [" + str(type(caller_token)) + "]")
return _get_header("Bearer", caller_token)
# Obtain token from MSI endpoint
if self.authentication_method == AuthenticationMethod.aad_msi:
token = self.get_token_from_msi()
return _get_header_from_dict(token)
refresh_token = None
if self.authentication_method == AuthenticationMethod.az_cli_profile:
stored_token = _get_azure_cli_auth_token()
if (
TokenResponseFields.REFRESH_TOKEN in stored_token
and TokenResponseFields._CLIENT_ID in stored_token
and TokenResponseFields._AUTHORITY in stored_token
):
self.client_id = stored_token[TokenResponseFields._CLIENT_ID]
self.username = stored_token[TokenResponseFields.USER_ID]
self.authority_uri = stored_token[TokenResponseFields._AUTHORITY]
refresh_token = stored_token[TokenResponseFields.REFRESH_TOKEN]
if self.auth_context is None:
self.auth_context = AuthenticationContext(self.authority_uri)
if refresh_token is not None:
token = self.auth_context.acquire_token_with_refresh_token(refresh_token, self.client_id, self.kusto_uri)
else:
token = self.auth_context.acquire_token(self.kusto_uri, self.username, self.client_id)
if token is not None:
expiration_date = dateutil.parser.parse(token[TokenResponseFields.EXPIRES_ON])
if expiration_date > datetime.now() + timedelta(minutes=1):
return _get_header_from_dict(token)
if TokenResponseFields.REFRESH_TOKEN in token:
token = self.auth_context.acquire_token_with_refresh_token(token[TokenResponseFields.REFRESH_TOKEN], self.client_id, self.kusto_uri)
if token is not None:
return _get_header_from_dict(token)
# obtain token from AAD
if self.authentication_method is AuthenticationMethod.aad_username_password:
token = self.auth_context.acquire_token_with_username_password(self.kusto_uri, self.username, self.password, self.client_id)
elif self.authentication_method is AuthenticationMethod.aad_application_key:
token = self.auth_context.acquire_token_with_client_credentials(self.kusto_uri, self.client_id, self.client_secret)
elif self.authentication_method is AuthenticationMethod.aad_device_login:
code = self.auth_context.acquire_user_code(self.kusto_uri, self.client_id)
print(code[OAuth2DeviceCodeResponseParameters.MESSAGE])
webbrowser.open(code[OAuth2DeviceCodeResponseParameters.VERIFICATION_URL])
token = self.auth_context.acquire_token_with_device_code(self.kusto_uri, code, self.client_id)
elif self.authentication_method is AuthenticationMethod.aad_application_certificate:
token = self.auth_context.acquire_token_with_client_certificate(self.kusto_uri, self.client_id, self.certificate, self.thumbprint)
else:
raise KustoClientError("Please choose authentication method from azure.kusto.data.security.AuthenticationMethod")
return _get_header_from_dict(token)
def get_token_from_msi(self) -> dict:
try:
if self.msi_auth_context is None:
# Create the MSI Authentication object, the first token is acquired implicitly
self.msi_auth_context = MSIAuthentication(**self.msi_params)
else:
# Acquire a fresh token
self.msi_auth_context.set_token()
except Exception as e:
raise KustoClientError("Failed to obtain MSI context for [" + str(self.msi_params) + "]\n" + str(e))
return self.msi_auth_context.token
def test_msi_vm(self):
# Test legacy MSI, with no MSI_ENDPOINT
json_payload = {
'token_type': "TokenType",
"access_token": "AccessToken"
}
httpretty.register_uri(httpretty.POST,
'http://localhost:666/oauth2/token',
body=json.dumps(json_payload),
content_type="application/json")
token_type, access_token, token_entry = get_msi_token("whatever", port=666)
assert token_type == "TokenType"
assert access_token == "AccessToken"
assert token_entry == json_payload
httpretty.register_uri(httpretty.POST,
'http://localhost:42/oauth2/token',
status=503,
content_type="application/json")
with self.assertRaises(HTTPError):
get_msi_token("whatever", port=42)
# Test MSI_ENDPOINT
json_payload = {
'token_type': "TokenType",
"access_token": "AccessToken"
}
httpretty.register_uri(httpretty.POST,
'http://random.org/yadadada',
body=json.dumps(json_payload),
content_type="application/json")
with mock.patch('os.environ', {'MSI_ENDPOINT': 'http://random.org/yadadada'}):
token_type, access_token, token_entry = get_msi_token("whatever")
assert token_type == "TokenType"
assert access_token == "AccessToken"
assert token_entry == json_payload
# Test MSIAuthentication with no MSI_ENDPOINT and no APPSETTING_WEBSITE_SITE_NAME is IMDS
json_payload = {
'token_type': "TokenTypeIMDS",
"access_token": "AccessToken"
}
httpretty.register_uri(httpretty.GET,
'http://169.254.169.254/metadata/identity/oauth2/token',
body=json.dumps(json_payload),
content_type="application/json")
credentials = MSIAuthentication()
credentials.set_token()
assert credentials.scheme == "TokenTypeIMDS"
assert credentials.token == json_payload
# Test MSIAuthentication with MSI_ENDPOINT and no APPSETTING_WEBSITE_SITE_NAME is MSI_ENDPOINT
json_payload = {
'token_type': "TokenTypeMSI_ENDPOINT",
"access_token": "AccessToken"
}
httpretty.register_uri(httpretty.POST,
'http://random.org/yadadada',
body=json.dumps(json_payload),
content_type="application/json")
with mock.patch('os.environ', {'MSI_ENDPOINT': 'http://random.org/yadadada'}):
credentials = MSIAuthentication()
credentials.set_token()
assert credentials.scheme == "TokenTypeMSI_ENDPOINT"
assert credentials.token == json_payload
# WebApp
json_payload = {
'token_type': "TokenTypeWebApp",
"access_token": "AccessToken"
}
httpretty.register_uri(httpretty.GET,
'http://127.0.0.1:41741/MSI/token/?resource=foo&api-version=2017-09-01',
body=json.dumps(json_payload),
content_type="application/json")
app_service_env = {
'APPSETTING_WEBSITE_SITE_NAME': 'Website name',
'MSI_ENDPOINT': 'http://127.0.0.1:41741/MSI/token',
'MSI_SECRET': '69418689F1E342DD946CB82994CDA3CB'
}
with mock.patch.dict('os.environ', app_service_env):
credentials = MSIAuthentication(resource="foo")
credentials.set_token()
assert credentials.scheme == "TokenTypeWebApp"
assert credentials.token == json_payload
