diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py index 0bf046e9573..354589a2fc8 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/custom.py +++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py @@ -533,8 +533,8 @@ def _build_service_principal(rbac_client, cli_ctx, name, url, client_secret): # always create application with 5 years expiration start_date = datetime.datetime.utcnow() end_date = start_date + relativedelta(years=5) - result = create_application(rbac_client.applications, name, url, [url], password=client_secret, - start_date=start_date, end_date=end_date) + result, aad_session_key = create_application(rbac_client.applications, name, url, [url], password=client_secret, + start_date=start_date, end_date=end_date) service_principal = result.app_id # pylint: disable=no-member for x in range(0, 10): hook.add(message='Creating service principal', value=0.1 * x, total_val=1.0) @@ -546,10 +546,10 @@ def _build_service_principal(rbac_client, cli_ctx, name, url, client_secret): logger.info(ex) time.sleep(2 + 2 * x) else: - return False + return False, aad_session_key hook.add(message='Finished service principal creation', value=1.0, total_val=1.0) logger.info('Finished service principal creation') - return service_principal + return service_principal, aad_session_key def _add_role_assignment(cli_ctx, role, service_principal_msi_id, is_service_principal=True, delay=2, scope=None): @@ -1313,7 +1313,8 @@ def create_application(client, display_name, homepage, identifier_uris, password_credentials=password_creds, required_resource_access=required_resource_accesses) try: - return client.create(app_create_param) + result = client.create(app_create_param, raw=True) + return result.output, result.response.headers["ocp-aad-session-key"] except GraphErrorException as ex: if 'insufficient privileges' in str(ex).lower(): link = 'https://docs.microsoft.com/azure/azure-resource-manager/resource-group-create-service-principal-portal' # pylint: disable=line-too-long @@ -1840,6 +1841,9 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value, # pylint: identity=identity ) + # Add AAD session key to header + custom_headers = {'Ocp-Aad-Session-Key': principal_obj.get("aad_session_key")} + # Due to SPN replication latency, we do a few retries here max_retry = 30 retry_exception = Exception(None) @@ -1867,7 +1871,9 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value, # pylint: result = sdk_no_wait(no_wait, client.create_or_update, resource_group_name=resource_group_name, - resource_name=name, parameters=mc) + resource_name=name, + parameters=mc, + custom_headers=custom_headers) return result except CloudError as ex: retry_exception = ex @@ -2896,6 +2902,7 @@ def _ensure_aks_service_principal(cli_ctx, dns_name_prefix=None, location=None, name=None): + aad_session_key = None # TODO: This really needs to be unit tested. rbac_client = get_graph_rbac_management_client(cli_ctx) if not service_principal: @@ -2905,7 +2912,7 @@ def _ensure_aks_service_principal(cli_ctx, salt = binascii.b2a_hex(os.urandom(3)).decode('utf-8') url = 'https://{}.{}.{}.cloudapp.azure.com'.format(salt, dns_name_prefix, location) - service_principal = _build_service_principal(rbac_client, cli_ctx, name, url, client_secret) + service_principal, aad_session_key = _build_service_principal(rbac_client, cli_ctx, name, url, client_secret) if not service_principal: raise CLIError('Could not create a service principal with the right permissions. ' 'Are you an Owner on this project?') @@ -2918,6 +2925,7 @@ def _ensure_aks_service_principal(cli_ctx, return { 'client_secret': client_secret, 'service_principal': service_principal, + 'aad_session_key': aad_session_key, } @@ -2962,12 +2970,12 @@ def _ensure_osa_aad(cli_ctx, required_resource_accesses=[required_osa_aad_access]) logger.info('Updated AAD: %s', aad_client_app_id) else: - result = create_application(client=rbac_client.applications, - display_name=name, - identifier_uris=[app_id_name], - homepage=app_id_name, - password=aad_client_app_secret, - required_resource_accesses=[required_osa_aad_access]) + result, _aad_session_key = create_application(client=rbac_client.applications, + display_name=name, + identifier_uris=[app_id_name], + homepage=app_id_name, + password=aad_client_app_secret, + required_resource_accesses=[required_osa_aad_access]) aad_client_app_id = result.app_id logger.info('Created an AAD: %s', aad_client_app_id) # Get the TenantID @@ -2998,7 +3006,7 @@ def _ensure_service_principal(cli_ctx, salt = binascii.b2a_hex(os.urandom(3)).decode('utf-8') url = 'https://{}.{}.{}.cloudapp.azure.com'.format(salt, dns_name_prefix, location) - service_principal = _build_service_principal(rbac_client, cli_ctx, name, url, client_secret) + service_principal, _aad_session_key = _build_service_principal(rbac_client, cli_ctx, name, url, client_secret) if not service_principal: raise CLIError('Could not create a service principal with the right permissions. ' 'Are you an Owner on this project?') diff --git a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_service_principals.py b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_service_principals.py index 7797c7f5c6c..9f7a7140aba 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_service_principals.py +++ b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_service_principals.py @@ -77,7 +77,9 @@ def test_build_service_principal(self): client = mock.MagicMock() client.service_principals = mock.Mock() client.applications = mock.Mock() - client.applications.create.return_value.app_id = app_id + result = mock.MagicMock() + result.output.app_id = app_id + client.applications.create.return_value = result client.applications.list.return_value = [] cli_ctx = mock.MagicMock() @@ -96,7 +98,7 @@ def test_build_service_principal(self): ] client.applications.list.assert_has_calls(expected_calls) # TODO better matcher here - client.applications.create.assert_called_with(mock.ANY) + client.applications.create.assert_called_with(mock.ANY, raw=True) if __name__ == '__main__':