def test_multitenant_authentication_not_allowed():
expected_tenant = "expected-tenant"
expected_token = "***"
def fake_check_output(command_line, **_):
match = re.search("--tenant (.*)", command_line[-1])
assert match is None or match[1] == expected_tenant
return json.dumps(
{
"expiresOn": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
"accessToken": expected_token,
"subscription": "some-guid",
"tenant": expected_token,
"tokenType": "Bearer",
}
)
credential = AzureCliCredential()
with mock.patch(CHECK_OUTPUT, fake_check_output):
token = credential.get_token("scope")
assert token.token == expected_token
with mock.patch.dict(
"os.environ", {EnvironmentVariables.AZURE_IDENTITY_DISABLE_MULTITENANTAUTH: "true"}
):
token = credential.get_token("scope", tenant_id="un" + expected_tenant)
assert token.token == expected_token
def _get_azcli_token(self, subscription: str = None) -> str:
"retrieve token from azcli login"
token = None
tenant = self._authority if subscription is None else None
self._current_authentication_method = self._current_authentication_method = AuthenticationMethod.azcli_login_subscription if subscription is not None else AuthenticationMethod.azcli_login
try:
from azure.identity import AzureCliCredential
try:
credential = AzureCliCredential()
access_token = credential.get_token(self._resource)
expires_datetime = datetime.fromtimestamp(
access_token.expires_on)
token = {
'accessToken': access_token.token,
'expiresOn':
expires_datetime.strftime("%Y-%m-%d %H:%M:%S.%f"),
'tokenType': 'Bearer',
}
except:
pass
except [ImportError, ModuleNotFoundError]:
raise AuthenticationError(
"Azure CLI authentication requires 'azure-cli-core' to be installed."
)
except:
pass
logger().debug(
f"_MyAadHelper::_get_azcli_token {'failed' if token is None else 'succeeded'} to get token - subscription: '{subscription}', tenant: '{tenant}'"
)
return token
def _get_token(self):
try:
credential = AzureCliCredential()
azureToken = credential.get_token(self.scope)
except Exception as ex:
raise RuntimeError(str(ex))
return azureToken.token
def _get_token_impl(self) -> dict:
try:
if self._az_auth_context is None:
self._az_auth_context = AzureCliCredential()
self._az_token = self._az_auth_context.get_token(self._kusto_uri)
return {TokenConstants.AZ_TOKEN_TYPE: TokenConstants.BEARER_TYPE, TokenConstants.AZ_ACCESS_TOKEN: self._az_token.token}
except Exception as e:
raise KustoClientError(
"Failed to obtain Az Cli token for '{0}'.\nPlease be sure AzCli version 2.3.0 and above is intalled.\n{1}".format(self._kusto_uri, e)
)
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)
self.credential = credential
if credential is None:
self.credential = AzureCliCredential()
self._policy = BearerTokenCredentialPolicy(self.credential, resource_id, **kwargs)
def _configure_resource_group(config):
# TODO: look at availability sets
# https://docs.microsoft.com/en-us/azure/virtual-machines/windows/tutorial-availability-sets
subscription_id = config["provider"].get("subscription_id")
if subscription_id is None:
subscription_id = get_cli_profile().get_subscription_id()
resource_client = ResourceManagementClient(AzureCliCredential(),
subscription_id)
config["provider"]["subscription_id"] = subscription_id
logger.info("Using subscription id: %s", subscription_id)
assert ("resource_group" in config["provider"]
), "Provider config must include resource_group field"
resource_group = config["provider"]["resource_group"]
assert (
"location"
in config["provider"]), "Provider config must include location field"
params = {"location": config["provider"]["location"]}
if "tags" in config["provider"]:
params["tags"] = config["provider"]["tags"]
logger.info("Creating/Updating Resource Group: %s", resource_group)
rg_create_or_update = get_azure_sdk_function(
client=resource_client.resource_groups,
function_name="create_or_update")
rg_create_or_update(resource_group_name=resource_group, parameters=params)
# load the template file
current_path = Path(__file__).parent
template_path = current_path.joinpath("azure-config-template.json")
with open(template_path, "r") as template_fp:
template = json.load(template_fp)
# choose a random subnet, skipping most common value of 0
random.seed(resource_group)
subnet_mask = "10.{}.0.0/16".format(random.randint(1, 254))
parameters = {
"properties": {
"mode": DeploymentMode.incremental,
"template": template,
"parameters": {
"subnet": {
"value": subnet_mask
}
},
}
}
create_or_update = get_azure_sdk_function(
client=resource_client.deployments, function_name="create_or_update")
create_or_update(
resource_group_name=resource_group,
deployment_name="ray-config",
parameters=parameters,
).wait()
return config
def test_not_logged_in():
"""When the CLI isn't logged in, the credential should raise CredentialUnavailableError"""
output = "ERROR: Please run 'az login' to setup account."
with mock.patch(CHECK_OUTPUT, raise_called_process_error(1, output)):
with pytest.raises(CredentialUnavailableError, match=NOT_LOGGED_IN):
AzureCliCredential().get_token("scope")
def test_not_logged_in():
"""When the CLI isn't logged in, the credential should raise an error containing the CLI's output"""
output = "ERROR: Please run 'az login' to setup account."
with mock.patch(CHECK_OUTPUT, raise_called_process_error(1, output)):
with pytest.raises(ClientAuthenticationError, match=output):
AzureCliCredential().get_token("scope")
def test_output(run_output):
credential = AzureCliCredential()
subscription_id = 'c65eaf3d-141a-4b75-9642-61a4c76ae763'
resource_client = ResourceManagementClient(credential, subscription_id)
assert resource_client.resource_groups.get(
'myResourceGroup').name == run_output['something']
print('Testing clean apply')
def ml_workspace_list(cmd, resource_group_name):
subscription_id = get_subscription_id(cmd.cli_ctx)
ml_client = MLClient(subscription_id=subscription_id,
resource_group_name=resource_group_name,
credential=AzureCliCredential())
return ml_client.workspaces.list()
def ml_model_create(cmd,
resource_group_name,
workspace_name,
name=None,
version=None,
file=None,
path=None,
params_override=[]):
subscription_id = get_subscription_id(cmd.cli_ctx)
ml_client = MLClient(
subscription_id=subscription_id,
resource_group_name=resource_group_name,
default_workspace_name=workspace_name,
credential=AzureCliCredential(),
)
if name is not None:
params_override.append({"name": name})
if version is not None:
params_override.append({"version": version})
if path is not None:
params_override.append({"asset_path": path})
try:
return ml_client.model.create_or_update(
file=file, params_override=params_override)
except Exception as err:
print_error_and_exit(str(err))
def ml_datastore_attach_blob(
cmd,
resource_group_name,
workspace_name,
account_name,
container_name,
name,
account_key=None,
sas_token=None,
protocol=None,
endpoint=None,
):
subscription_id = get_subscription_id(cmd.cli_ctx)
ml_client = MLClient(
subscription_id=subscription_id,
resource_group_name=resource_group_name,
default_workspace_name=workspace_name,
credential=AzureCliCredential(),
)
return ml_client.datastores.attach_azure_blob_storage(
name,
container_name,
account_name,
account_key=account_key,
sas_token=sas_token,
protocol=protocol,
endpoint=endpoint,
)
class AzCliTokenProvider(CloudInfoTokenProvider):
"""AzCli Token Provider obtains a refresh token from the AzCli cache and uses it to authenticate with MSAL"""
def __init__(self, kusto_uri: str, is_async: bool = False):
super().__init__(kusto_uri, is_async)
self._az_auth_context = None
self._az_auth_context_async = None
self._az_token = None
@staticmethod
def name() -> str:
return "AzCliTokenProvider"
def _context_impl(self) -> dict:
return {"authority:": self.name()}
def _init_impl(self):
pass
def _get_token_impl(self) -> Optional[dict]:
try:
if self._az_auth_context is None:
self._az_auth_context = AzureCliCredential()
self._az_token = self._az_auth_context.get_token(self._scopes[0])
return {
TokenConstants.AZ_TOKEN_TYPE: TokenConstants.BEARER_TYPE,
TokenConstants.AZ_ACCESS_TOKEN: self._az_token.token
}
except Exception as e:
raise KustoClientError(
"Failed to obtain Az Cli token for '{0}'.\nPlease be sure AzCli version 2.3.0 and above is intalled.\n{1}"
.format(self._kusto_uri, e))
async def _get_token_impl_async(self) -> Optional[dict]:
try:
if self._az_auth_context_async is None:
self._az_auth_context_async = AsyncAzureCliCredential()
self._az_token = await self._az_auth_context_async.get_token(
self._scopes[0])
return {
TokenConstants.AZ_TOKEN_TYPE: TokenConstants.BEARER_TYPE,
TokenConstants.AZ_ACCESS_TOKEN: self._az_token.token
}
except Exception as e:
raise KustoClientError(
"Failed to obtain Az Cli token for '{0}'.\nPlease be sure AzCli version 2.3.0 and above is installed.\n{1}"
.format(self._kusto_uri, e))
def _get_token_from_cache_impl(self) -> Optional[dict]:
if self._az_token is not None:
# A token is considered valid if it is due to expire in no less than 10 minutes
cur_time = time.time()
if (self._az_token.expires_on - 600) > cur_time:
return {
TokenConstants.MSAL_TOKEN_TYPE: TokenConstants.BEARER_TYPE,
TokenConstants.MSAL_ACCESS_TOKEN: self._az_token.token
}
return None
def test_cli_not_installed_linux():
"""The credential should raise CredentialUnavailableError when the CLI isn't installed"""
output = "/bin/sh: 1: az: not found"
with mock.patch(CHECK_OUTPUT, raise_called_process_error(127, output)):
with pytest.raises(CredentialUnavailableError, match=CLI_NOT_FOUND):
AzureCliCredential().get_token("scope")
def test_unexpected_error():
"""When the CLI returns an unexpected error, the credential should raise an error containing the CLI's output"""
output = "something went wrong"
with mock.patch(CHECK_OUTPUT, raise_called_process_error(42, output)):
with pytest.raises(ClientAuthenticationError, match=output):
AzureCliCredential().get_token("scope")
def test_cli_not_installed_windows():
"""The credential should raise CredentialUnavailableError when the CLI isn't installed"""
output = "'az' is not recognized as an internal or external command, operable program or batch file."
with mock.patch(CHECK_OUTPUT, raise_called_process_error(1, output)):
with pytest.raises(CredentialUnavailableError, match=CLI_NOT_FOUND):
AzureCliCredential().get_token("scope")
def azure_credentials(azure_cfg: AzureCfg, pytestconfig):
if pytestconfig.getoption('local'):
return AzureCliCredential()
else:
return ClientSecretCredential(client_id=azure_cfg.client_id,
client_secret=azure_cfg.client_secret,
tenant_id=azure_cfg.tenant_id)
def test_get_token():
"""The credential should parse the CLI's output to an AccessToken"""
access_token = "access token"
expected_expires_on = 1602015811
successful_output = json.dumps({
"expiresOn":
datetime.fromtimestamp(expected_expires_on).strftime(
"%Y-%m-%d %H:%M:%S.%f"),
"accessToken":
access_token,
"subscription":
"some-guid",
"tenant":
"some-guid",
"tokenType":
"Bearer",
})
with mock.patch(CHECK_OUTPUT, mock.Mock(return_value=successful_output)):
token = AzureCliCredential().get_token("scope")
assert token.token == access_token
assert type(token.expires_on) == int
assert token.expires_on == expected_expires_on
def create_resource_group(subID):
credential = AzureCliCredential()
client = ResourceManagementClient(credential, subID)
client.resource_groups.create_or_update(resource_group_name=name,
parameters={"location": location})
def test_get_token():
"""The credential should parse the CLI's output to an AccessToken"""
access_token = "access token"
valid_seconds = 42
successful_output = json.dumps({
# expiresOn is a naive datetime representing valid_seconds from the epoch
"expiresOn":
datetime.fromtimestamp(valid_seconds).strftime("%Y-%m-%d %H:%M:%S.%f"),
"accessToken":
access_token,
"subscription":
"some-guid",
"tenant":
"some-guid",
"tokenType":
"Bearer",
})
with mock.patch(CHECK_OUTPUT, mock.Mock(return_value=successful_output)):
token = AzureCliCredential().get_token("scope")
assert token.token == access_token
assert type(token.expires_on) == int
assert token.expires_on == valid_seconds
def add_log_export(self) -> None:
if not self.export_appinsights:
logger.info("not exporting appinsights")
return
container_name = "app-insights"
logger.info("adding appinsight log export")
account_name = self.results["deploy"]["func_name"]["value"]
key = self.results["deploy"]["func_key"]["value"]
account_url = "https://%s.blob.core.windows.net" % account_name
client = BlobServiceClient(account_url, credential=key)
if container_name not in [x["name"] for x in client.list_containers()]:
client.create_container(container_name)
expiry = datetime.utcnow() + timedelta(days=2 * 365)
# NOTE: as this is a long-lived SAS url, it should not be logged and only
# used in the the later-on export_configurations.create() call
sas = generate_container_sas(
account_name,
container_name,
account_key=key,
permission=ContainerSasPermissions(write=True),
expiry=expiry,
)
url = "%s/%s?%s" % (account_url, container_name, sas)
record_types = (
"Requests, Event, Exceptions, Metrics, PageViews, "
"PageViewPerformance, Rdd, PerformanceCounters, Availability")
req = ApplicationInsightsComponentExportRequest(
record_types=record_types,
destination_type="Blob",
is_enabled="true",
destination_address=url,
)
credential = AzureCliCredential()
app_insight_client = ApplicationInsightsManagementClient(
credential,
subscription_id=self.get_subscription_id(),
)
to_delete = []
for entry in app_insight_client.export_configurations.list(
self.resource_group, self.application_name):
if (entry.storage_name == account_name
and entry.container_name == container_name):
to_delete.append(entry.export_id)
for export_id in to_delete:
logger.info("replacing existing export: %s", export_id)
app_insight_client.export_configurations.delete(
self.resource_group, self.application_name, export_id)
app_insight_client.export_configurations.create(
self.resource_group, self.application_name, req)
def _default_chained_credentials(self) -> ChainedTokenCredential:
managed_identity = ManagedIdentityCredential()
azure_cli = AzureCliCredential()
environment = EnvironmentCredential()
shared_token_cache = SharedTokenCacheCredential()
interactive_browser = InteractiveBrowserCredential()
return ChainedTokenCredential(managed_identity, azure_cli, environment, shared_token_cache, interactive_browser)
def get_azure_token():
credential_chain = ChainedTokenCredential(ManagedIdentityCredential(), AzureCliCredential())
try:
token = credential_chain.get_token("https://monitoring.azure.com//.default")
return token.token
except Exception as e:
logging.exception(f"Failed to retrieve Azure token. Reason is {type(e).__name__} {e}")
return None
def test_timeout():
"""The credential should raise CredentialUnavailableError when the subprocess times out"""
from subprocess import TimeoutExpired
with mock.patch(CHECK_OUTPUT, mock.Mock(side_effect=TimeoutExpired("", 42))):
with pytest.raises(CredentialUnavailableError):
AzureCliCredential().get_token("scope")
def test_parsing_error_does_not_expose_token(output):
"""Errors during CLI output parsing shouldn't expose access tokens in that output"""
with mock.patch(CHECK_OUTPUT, mock.Mock(return_value=output)):
with pytest.raises(ClientAuthenticationError) as ex:
AzureCliCredential().get_token("scope")
assert "secret value" not in str(ex.value)
assert "secret value" not in repr(ex.value)
def adf_client(adf_config):
"""Creates an DataFactoryManagementClient object"""
if adf_config["AZ_SERVICE_PRINCIPAL_ID"] is None:
credentials = AzureCliCredential()
else:
credentials = ClientSecretCredential(client_id=adf_config["AZ_SERVICE_PRINCIPAL_ID"],
client_secret=adf_config["AZ_SERVICE_PRINCIPAL_SECRET"],
tenant_id=adf_config["AZ_SERVICE_PRINCIPAL_TENANT_ID"])
return DataFactoryManagementClient(credentials, adf_config["AZ_SUBSCRIPTION_ID"])
def test_subprocess_error_does_not_expose_token(output):
"""Errors from the subprocess shouldn't expose access tokens in CLI output"""
with mock.patch(CHECK_OUTPUT, raise_called_process_error(1, output=output)):
with pytest.raises(ClientAuthenticationError) as ex:
AzureCliCredential().get_token("scope")
assert "secret value" not in str(ex.value)
assert "secret value" not in repr(ex.value)
def ml_compute_show(cmd, resource_group_name, workspace_name, name):
subscription_id = get_subscription_id(cmd.cli_ctx)
ml_client = MLClient(
subscription_id=subscription_id,
resource_group_name=resource_group_name,
default_workspace_name=workspace_name,
credential=AzureCliCredential(),
)
return ml_client.computes.get(name=name)
def ml_data_delete(cmd, resource_group_name, workspace_name, name, version):
subscription_id = get_subscription_id(cmd.cli_ctx)
ml_client = MLClient(
subscription_id=subscription_id,
resource_group_name=resource_group_name,
default_workspace_name=workspace_name,
credential=AzureCliCredential(),
)
return ml_client.data.delete(name=name, version=version)
def ml_data_list(cmd, resource_group_name, workspace_name, name=None):
subscription_id = get_subscription_id(cmd.cli_ctx)
ml_client = MLClient(
subscription_id=subscription_id,
resource_group_name=resource_group_name,
default_workspace_name=workspace_name,
credential=AzureCliCredential(),
)
return ml_client.data.list(name=name)
