/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.ambari.server.controller;

import static org.apache.ambari.server.controller.KerberosHelper.DIRECTIVE_COMPONENTS;
import static org.apache.ambari.server.controller.KerberosHelper.DIRECTIVE_HOSTS;
import static org.easymock.EasyMock.anyLong;
import static org.easymock.EasyMock.anyObject;
import static org.easymock.EasyMock.anyString;
import static org.easymock.EasyMock.capture;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.expectLastCall;
import static org.easymock.EasyMock.getCurrentArguments;
import static org.easymock.EasyMock.isNull;
import static org.easymock.EasyMock.newCapture;
import static org.easymock.EasyMock.replay;
import static org.easymock.EasyMock.reset;
import static org.easymock.EasyMock.verify;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;

import java.io.File;
import java.lang.reflect.Method;
import java.net.InetAddress;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.TimeUnit;

import jakarta.persistence.EntityManager;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.Role;
import org.apache.ambari.server.actionmanager.ActionDBAccessor;
import org.apache.ambari.server.actionmanager.ActionDBAccessorImpl;
import org.apache.ambari.server.actionmanager.ActionManager;
import org.apache.ambari.server.actionmanager.HostRoleCommandFactory;
import org.apache.ambari.server.actionmanager.HostRoleCommandFactoryImpl;
import org.apache.ambari.server.actionmanager.Stage;
import org.apache.ambari.server.actionmanager.StageFactory;
import org.apache.ambari.server.api.services.AmbariMetaInfo;
import org.apache.ambari.server.api.services.stackadvisor.StackAdvisorHelper;
import org.apache.ambari.server.api.services.stackadvisor.StackAdvisorRequest;
import org.apache.ambari.server.api.services.stackadvisor.recommendations.RecommendationResponse;
import org.apache.ambari.server.audit.AuditLogger;
import org.apache.ambari.server.configuration.Configuration;
import org.apache.ambari.server.controller.internal.RequestStageContainer;
import org.apache.ambari.server.controller.spi.ClusterController;
import org.apache.ambari.server.controller.utilities.KerberosChecker;
import org.apache.ambari.server.hooks.HookService;
import org.apache.ambari.server.hooks.users.UserHookService;
import org.apache.ambari.server.ldap.service.LdapFacade;
import org.apache.ambari.server.metadata.CachedRoleCommandOrderProvider;
import org.apache.ambari.server.metadata.RoleCommandOrder;
import org.apache.ambari.server.metadata.RoleCommandOrderProvider;
import org.apache.ambari.server.mpack.MpackManagerFactory;
import org.apache.ambari.server.orm.DBAccessor;
import org.apache.ambari.server.orm.dao.ArtifactDAO;
import org.apache.ambari.server.orm.dao.HostRoleCommandDAO;
import org.apache.ambari.server.orm.dao.KerberosKeytabPrincipalDAO;
import org.apache.ambari.server.orm.dao.KerberosKeytabPrincipalDAO.KeytabPrincipalFindOrCreateResult;
import org.apache.ambari.server.orm.dao.KerberosPrincipalDAO;
import org.apache.ambari.server.orm.entities.KerberosKeytabPrincipalEntity;
import org.apache.ambari.server.scheduler.ExecutionScheduler;
import org.apache.ambari.server.scheduler.ExecutionSchedulerImpl;
import org.apache.ambari.server.security.SecurityHelper;
import org.apache.ambari.server.security.credential.PrincipalKeyCredential;
import org.apache.ambari.server.security.encryption.CredentialStoreService;
import org.apache.ambari.server.security.encryption.CredentialStoreServiceImpl;
import org.apache.ambari.server.security.encryption.CredentialStoreType;
import org.apache.ambari.server.serveraction.ActionLog;
import org.apache.ambari.server.serveraction.kerberos.ConfigureAmbariIdentitiesServerAction;
import org.apache.ambari.server.serveraction.kerberos.CreateKeytabFilesServerAction;
import org.apache.ambari.server.serveraction.kerberos.CreatePrincipalsServerAction;
import org.apache.ambari.server.serveraction.kerberos.KDCType;
import org.apache.ambari.server.serveraction.kerberos.KerberosConfigDataFileWriterFactory;
import org.apache.ambari.server.serveraction.kerberos.KerberosInvalidConfigurationException;
import org.apache.ambari.server.serveraction.kerberos.KerberosMissingAdminCredentialsException;
import org.apache.ambari.server.serveraction.kerberos.KerberosOperationException;
import org.apache.ambari.server.serveraction.kerberos.KerberosOperationHandler;
import org.apache.ambari.server.serveraction.kerberos.KerberosOperationHandlerFactory;
import org.apache.ambari.server.serveraction.kerberos.PreconfigureServiceType;
import org.apache.ambari.server.serveraction.kerberos.stageutils.ResolvedKerberosPrincipal;
import org.apache.ambari.server.stack.StackManagerFactory;
import org.apache.ambari.server.state.Cluster;
import org.apache.ambari.server.state.Clusters;
import org.apache.ambari.server.state.ComponentInfo;
import org.apache.ambari.server.state.Config;
import org.apache.ambari.server.state.ConfigHelper;
import org.apache.ambari.server.state.Host;
import org.apache.ambari.server.state.HostState;
import org.apache.ambari.server.state.PropertyInfo;
import org.apache.ambari.server.state.SecurityType;
import org.apache.ambari.server.state.Service;
import org.apache.ambari.server.state.ServiceComponent;
import org.apache.ambari.server.state.ServiceComponentHost;
import org.apache.ambari.server.state.ServiceComponentHostFactory;
import org.apache.ambari.server.state.ServiceInfo;
import org.apache.ambari.server.state.StackId;
import org.apache.ambari.server.state.State;
import org.apache.ambari.server.state.kerberos.KerberosComponentDescriptor;
import org.apache.ambari.server.state.kerberos.KerberosConfigurationDescriptor;
import org.apache.ambari.server.state.kerberos.KerberosDescriptor;
import org.apache.ambari.server.state.kerberos.KerberosDescriptorFactory;
import org.apache.ambari.server.state.kerberos.KerberosIdentityDescriptor;
import org.apache.ambari.server.state.kerberos.KerberosKeytabDescriptor;
import org.apache.ambari.server.state.kerberos.KerberosPrincipalDescriptor;
import org.apache.ambari.server.state.kerberos.KerberosPrincipalType;
import org.apache.ambari.server.state.kerberos.KerberosServiceDescriptor;
import org.apache.ambari.server.state.stack.OsFamily;
import org.apache.ambari.server.testutils.PartialNiceMockBinder;
import org.apache.ambari.server.topology.PersistedState;
import org.apache.ambari.server.topology.PersistedStateImpl;
import org.apache.ambari.server.topology.TopologyManager;
import org.apache.ambari.server.utils.StageUtils;
import org.apache.directory.server.kerberos.shared.keytab.Keytab;
import org.easymock.Capture;
import org.easymock.CaptureType;
import org.easymock.EasyMock;
import org.easymock.EasyMockSupport;
import org.easymock.IAnswer;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.springframework.security.crypto.password.PasswordEncoder;
import org.springframework.security.crypto.password.StandardPasswordEncoder;

import com.google.inject.AbstractModule;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.persist.jpa.AmbariJpaPersistService;


@SuppressWarnings("unchecked")
public class KerberosHelperTest extends EasyMockSupport {

  private static Injector injector;
  private final ClusterController clusterController = createStrictMock(ClusterController.class);
  private final KerberosDescriptorFactory kerberosDescriptorFactory = createStrictMock(KerberosDescriptorFactory.class);
  private final KerberosConfigDataFileWriterFactory kerberosConfigDataFileWriterFactory = createStrictMock(KerberosConfigDataFileWriterFactory.class);
  private final AmbariMetaInfo metaInfo = createMock(AmbariMetaInfo.class);
  private final TopologyManager topologyManager = createMock(TopologyManager.class);
  private final Configuration configuration = createMock(Configuration.class);
  private final AmbariCustomCommandExecutionHelper customCommandExecutionHelperMock = createNiceMock(AmbariCustomCommandExecutionHelper.class);

  @Rule
  public TemporaryFolder temporaryFolder = new TemporaryFolder();

  @Before
  public void setUp() throws Exception {
    reset(clusterController);
    reset(metaInfo);

    final KerberosOperationHandlerFactory kerberosOperationHandlerFactory = createMock(KerberosOperationHandlerFactory.class);

    expect(kerberosOperationHandlerFactory.getKerberosOperationHandler(KDCType.NONE))
        .andReturn(null)
        .anyTimes();

    expect(kerberosOperationHandlerFactory.getKerberosOperationHandler(KDCType.MIT_KDC))
        .andReturn(new KerberosOperationHandler() {
          @Override
          public void open(PrincipalKeyCredential administratorCredentials, String defaultRealm, Map<String, String> kerberosConfiguration) throws KerberosOperationException {
            setAdministratorCredential(administratorCredentials);
            setDefaultRealm(defaultRealm);
            setOpen(true);
          }

          @Override
          public void close() throws KerberosOperationException {

          }

          @Override
          public boolean principalExists(String principal, boolean service) throws KerberosOperationException {
            return "principal".equals(principal);
          }

          @Override
          public Integer createPrincipal(String principal, String password, boolean service) throws KerberosOperationException {
            return null;
          }

          @Override
          public Integer setPrincipalPassword(String principal, String password, boolean service) throws KerberosOperationException {
            return null;
          }

          @Override
          public boolean removePrincipal(String principal, boolean service) throws KerberosOperationException {
            return false;
          }

          @Override
          public boolean createKeytabFile(Keytab keytab, File destinationKeytabFile) throws KerberosOperationException {
            return true;
          }
        })
        .anyTimes();

    Method methodGetConfiguredTemporaryDirectory = KerberosHelperImpl.class.getDeclaredMethod("getConfiguredTemporaryDirectory");

    final KerberosHelperImpl kerberosHelper = createMockBuilder(KerberosHelperImpl.class)
        .addMockedMethod(methodGetConfiguredTemporaryDirectory)
        .createMock();

    expect(kerberosHelper.getConfiguredTemporaryDirectory()).andReturn(temporaryFolder.getRoot()).anyTimes();

    injector = Guice.createInjector(new AbstractModule() {

      @Override
      protected void configure() {
        PartialNiceMockBinder.newBuilder().addActionDBAccessorConfigsBindings().addFactoriesInstallBinding()
            .addPasswordEncryptorBindings().build().configure(binder());

        bind(AmbariJpaPersistService.class).toInstance(createNiceMock(AmbariJpaPersistService.class));
        bind(ActionDBAccessor.class).to(ActionDBAccessorImpl.class);
        bind(ExecutionScheduler.class).to(ExecutionSchedulerImpl.class);
        bind(AbstractRootServiceResponseFactory.class).to(RootServiceResponseFactory.class);
        bind(ServiceComponentHostFactory.class).toInstance(createNiceMock(ServiceComponentHostFactory.class));
        bind(PasswordEncoder.class).toInstance(new StandardPasswordEncoder());
        bind(HookService.class).to(UserHookService.class);
        bind(PersistedState.class).to(PersistedStateImpl.class);
        bind(EntityManager.class).toInstance(createNiceMock(EntityManager.class));
        bind(DBAccessor.class).toInstance(createNiceMock(DBAccessor.class));
        bind(SecurityHelper.class).toInstance(createNiceMock(SecurityHelper.class));
        bind(OsFamily.class).toInstance(createNiceMock(OsFamily.class));
        bind(AmbariCustomCommandExecutionHelper.class).toInstance(customCommandExecutionHelperMock);
        bind(AmbariManagementController.class).toInstance(createNiceMock(AmbariManagementController.class));
        bind(AmbariMetaInfo.class).toInstance(metaInfo);
        bind(ActionManager.class).toInstance(createNiceMock(ActionManager.class));
        bind(StageFactory.class).toInstance(createNiceMock(StageFactory.class));
        bind(Clusters.class).toInstance(createNiceMock(Clusters.class));
        bind(ConfigHelper.class).toInstance(createNiceMock(ConfigHelper.class));
        bind(KerberosOperationHandlerFactory.class).toInstance(kerberosOperationHandlerFactory);
        bind(ClusterController.class).toInstance(clusterController);
        bind(KerberosDescriptorFactory.class).toInstance(kerberosDescriptorFactory);
        bind(KerberosConfigDataFileWriterFactory.class).toInstance(kerberosConfigDataFileWriterFactory);
        bind(StackManagerFactory.class).toInstance(createNiceMock(StackManagerFactory.class));
        bind(KerberosHelper.class).toInstance(kerberosHelper);
        bind(CredentialStoreService.class).to(CredentialStoreServiceImpl.class);
        bind(CreatePrincipalsServerAction.class).toInstance(createMock(CreatePrincipalsServerAction.class));
        bind(CreateKeytabFilesServerAction.class).toInstance(createMock(CreateKeytabFilesServerAction.class));
        bind(ConfigureAmbariIdentitiesServerAction.class).toInstance(createMock(ConfigureAmbariIdentitiesServerAction.class));
        bind(StackAdvisorHelper.class).toInstance(createMock(StackAdvisorHelper.class));
        bind(HostRoleCommandDAO.class).toInstance(createNiceMock(HostRoleCommandDAO.class));
        bind(AuditLogger.class).toInstance(createNiceMock(AuditLogger.class));
        bind(ArtifactDAO.class).toInstance(createNiceMock(ArtifactDAO.class));
        bind(KerberosPrincipalDAO.class).toInstance(createNiceMock(KerberosPrincipalDAO.class));
        bind(KerberosKeytabPrincipalDAO.class).toInstance(createNiceMock(KerberosKeytabPrincipalDAO.class));
        bind(RoleCommandOrderProvider.class).to(CachedRoleCommandOrderProvider.class);
        bind(HostRoleCommandFactory.class).to(HostRoleCommandFactoryImpl.class);
        bind(MpackManagerFactory.class).toInstance(createNiceMock(MpackManagerFactory.class));
        bind(LdapFacade.class).toInstance(createNiceMock(LdapFacade.class));
        requestStaticInjection(KerberosChecker.class);
      }
    });

    //todo: StageUtils shouldn't be called for this test
    StageUtils.setTopologyManager(topologyManager);
    expect(topologyManager.getPendingHostComponents()).andReturn(
        Collections.emptyMap()).anyTimes();

    StageUtils.setConfiguration(configuration);
    expect(configuration.getApiSSLAuthentication()).andReturn(false).anyTimes();
    expect(configuration.getClientApiPort()).andReturn(8080).anyTimes();
    expect(configuration.getServerTempDir()).andReturn(temporaryFolder.getRoot().getAbsolutePath()).anyTimes();

    CredentialStoreService credentialStoreService = injector.getInstance(CredentialStoreService.class);
    if (!credentialStoreService.isInitialized(CredentialStoreType.TEMPORARY)) {
      ((CredentialStoreServiceImpl) credentialStoreService).initializeTemporaryCredentialStore(10, TimeUnit.MINUTES, false);
    }
  }

  @After
  public void tearDown() throws Exception {

  }

  @Test(expected = AmbariException.class)
  public void testMissingClusterEnv() throws Exception {
    KerberosHelper kerberosHelper = injector.getInstance(KerberosHelper.class);

    Cluster cluster = createNiceMock(Cluster.class);
    RequestStageContainer requestStageContainer = createNiceMock(RequestStageContainer.class);

    replayAll();
    kerberosHelper.toggleKerberos(cluster, SecurityType.KERBEROS, requestStageContainer, true);
    verifyAll();
  }

  @Test(expected = AmbariException.class)
  public void testMissingKrb5Conf() throws Exception {
    KerberosHelper kerberosHelper = injector.getInstance(KerberosHelper.class);

    final Map<String, String> kerberosEnvProperties = createMock(Map.class);
    expect(kerberosEnvProperties.get("ldap_url")).andReturn("").once();
    expect(kerberosEnvProperties.get("container_dn")).andReturn("").once();

    final Config kerberosEnvConfig = createMock(Config.class);
    expect(kerberosEnvConfig.getProperties()).andReturn(kerberosEnvProperties).once();

    final Cluster cluster = createNiceMock(Cluster.class);
    expect(cluster.getDesiredConfigByType("kerberos-env")).andReturn(kerberosEnvConfig).once();

    replayAll();
    kerberosHelper.toggleKerberos(cluster, SecurityType.KERBEROS, null, true);
    verifyAll();
  }

  @Test(expected = AmbariException.class)
  public void testMissingKerberosEnvConf() throws Exception {
    KerberosHelper kerberosHelper = injector.getInstance(KerberosHelper.class);

    final Map<String, String> kerberosEnvProperties = createMock(Map.class);
    expect(kerberosEnvProperties.get(KerberosHelper.DEFAULT_REALM)).andReturn("EXAMPLE.COM").once();
    expect(kerberosEnvProperties.get("kdc_hosts")).andReturn("10.0.100.1").once();

    final Map<String, String> krb5ConfProperties = createMock(Map.class);
    expect(krb5ConfProperties.get("kadmin_host")).andReturn("10.0.100.1").once();

    final Config krb5ConfConfig = createMock(Config.class);
    expect(krb5ConfConfig.getProperties()).andReturn(krb5ConfProperties).once();

    final Cluster cluster = createNiceMock(Cluster.class);
    expect(cluster.getDesiredConfigByType("krb5-conf")).andReturn(krb5ConfConfig).once();

    replayAll();
    kerberosHelper.toggleKerberos(cluster, SecurityType.KERBEROS, null, true);
    verifyAll();
  }

  @Test
  public void testEnableKerberos() throws Exception {
    testEnableKerberos(new PrincipalKeyCredential("principal", "password"), "mit-kdc", "true");
  }

  @Test
  public void testEnableKerberos_ManageIdentitiesFalseKdcNone() throws Exception {
    testEnableKerberos(new PrincipalKeyCredential("principal", "password"), "none", "false");
  }

  @Test(expected = AmbariException.class)
  public void testEnableKerberos_ManageIdentitiesTrueKdcNone() throws Exception {
    testEnableKerberos(new PrincipalKeyCredential("principal", "password"), "none", "true");
  }

  @Test(expected = KerberosInvalidConfigurationException.class)
  public void testEnableKerberos_ManageIdentitiesTrueKdcNull() throws Exception {
    testEnableKerberos(new PrincipalKeyCredential("principal", "password"), null, "true");
  }

  @Test(expected = KerberosMissingAdminCredentialsException.class)
  public void testEnableKerberosMissingCredentials() throws Exception {
    try {
      testEnableKerberos(null, "mit-kdc", "true");
    } catch (IllegalArgumentException e) {
      Assert.assertTrue(e.getMessage().startsWith("Missing KDC administrator credentials"));
      throw e;
    }
  }

  @Test(expected = KerberosMissingAdminCredentialsException.class)
  public void testEnableKerberosInvalidCredentials() throws Exception {
    try {
      testEnableKerberos(new PrincipalKeyCredential("invalid_principal", "password"), "mit-kdc", "true");
    } catch (IllegalArgumentException e) {
      Assert.assertTrue(e.getMessage().startsWith("Invalid KDC administrator credentials"));
      throw e;
    }
  }

  @Test
  public void testEnableKerberos_GetKerberosDescriptorFromCluster() throws Exception {
    testEnableKerberos(new PrincipalKeyCredential("principal", "password"), "mit-kdc", "true");
  }

  @Test
  public void testEnableKerberos_GetKerberosDescriptorFromStack() throws Exception {
    testEnableKerberos(new PrincipalKeyCredential("principal", "password"), "mit-kdc", "true");
  }

  @Test
  public void testEnsureIdentities() throws Exception {
    testEnsureIdentities(new PrincipalKeyCredential("principal", "password"), null);
  }

  @Test(expected = KerberosMissingAdminCredentialsException.class)
  public void testEnsureIdentitiesMissingCredentials() throws Exception {
    try {
      testEnsureIdentities(null, null);
    } catch (IllegalArgumentException e) {
      Assert.assertTrue(e.getMessage().startsWith("Missing KDC administrator credentials"));
      throw e;
    }
  }

  @Test(expected = KerberosMissingAdminCredentialsException.class)
  public void testEnsureIdentitiesInvalidCredentials() throws Exception {
    try {
      testEnsureIdentities(new PrincipalKeyCredential("invalid_principal", "password"), null);
    } catch (IllegalArgumentException e) {
      Assert.assertTrue(e.getMessage().startsWith("Invalid KDC administrator credentials"));
      throw e;
    }
  }

  @Test
  public void testEnsureIdentities_FilteredHosts() throws Exception {
    testEnsureIdentities(new PrincipalKeyCredential("principal", "password"), Collections.singleton("hostA"));
  }

  @Test
  public void testDeleteIdentities() throws Exception {
    testDeleteIdentities(new PrincipalKeyCredential("principal", "password"));
  }

  @Test(expected = KerberosMissingAdminCredentialsException.class)
  public void testDeleteIdentitiesMissingCredentials() throws Exception {
    try {
      testDeleteIdentities(null);
    } catch (IllegalArgumentException e) {
      Assert.assertTrue(e.getMessage().startsWith("Missing KDC administrator credentials"));
      throw e;
    }
  }

  @Test(expected = KerberosMissingAdminCredentialsException.class)
  public void testDeleteIdentitiesInvalidCredentials() throws Exception {
    try {
      testDeleteIdentities(new PrincipalKeyCredential("invalid_principal", "password"));
    } catch (IllegalArgumentException e) {
      Assert.assertTrue(e.getMessage().startsWith("Invalid KDC administrator credentials"));
      throw e;
    }
  }

  @Test
  public void testExecuteCustomOperationsInvalidOperation() throws Exception {
    KerberosHelper kerberosHelper = injector.getInstance(KerberosHelper.class);

    final Cluster cluster = createNiceMock(Cluster.class);

    try {
      kerberosHelper.executeCustomOperations(cluster,
          Collections.singletonMap("invalid_operation", "false"), null, true);
    } catch (Throwable t) {
      Assert.fail("Exception should not have been thrown");
    }
  }

  @Test(expected = AmbariException.class)
  public void testRegenerateKeytabsInvalidValue() throws Exception {
    KerberosHelper kerberosHelper = injector.getInstance(KerberosHelper.class);

    final Cluster cluster = createNiceMock(Cluster.class);

    kerberosHelper.executeCustomOperations(cluster,
        Collections.singletonMap(KerberosHelper.DIRECTIVE_REGENERATE_KEYTABS, "false"), null, true);
    Assert.fail("AmbariException should have failed");
  }

  @Test
  public void testRegenerateKeytabsValidateRequestStageContainer() throws Exception {
    testRegenerateKeytabs(new PrincipalKeyCredential("principal", "password"), true, false);
  }

  @Test
  public void testRegenerateKeytabsValidateSkipInvalidHost() throws Exception {
    testRegenerateKeytabs(new PrincipalKeyCredential("principal", "password"), true, true);
  }

  @Test
  public void testRegenerateKeytabs() throws Exception {
    testRegenerateKeytabs(new PrincipalKeyCredential("principal", "password"), false, false);
  }

  /**
   * Tests that when regenerating keytabs for an upgrade, that the retry allowed
   * boolean is set on the tasks created.
   *
   * @throws Exception
   */
  @Test
  public void testRegenerateKeytabsWithRetryAllowed() throws Exception {
    Capture<ActionExecutionContext> captureContext = Capture.newInstance();
    customCommandExecutionHelperMock.addExecutionCommandsToStage(capture(captureContext),
        anyObject(Stage.class), anyObject(), eq(null));

    expectLastCall().atLeastOnce();

    Map<String, String> requestMap = new HashMap<>();
    requestMap.put(KerberosHelper.DIRECTIVE_REGENERATE_KEYTABS, "true");
    requestMap.put(KerberosHelper.ALLOW_RETRY, "true");

    RequestStageContainer requestStageContainer = testRegenerateKeytabs(
        new PrincipalKeyCredential("principal", "password"), requestMap, false, false);

    assertNotNull(requestStageContainer);

    ActionExecutionContext capturedContext = captureContext.getValue();
    assertTrue(capturedContext.isRetryAllowed());
  }

  @Test
  public void testDisableKerberos() throws Exception {
    testDisableKerberos(new PrincipalKeyCredential("principal", "password"));
  }

  @Test
  public void testCreateTestIdentity_ManageIdentitiesDefault() throws Exception {
    testCreateTestIdentity(new PrincipalKeyCredential("principal", "password"), null);
  }

  @Test
  public void testCreateTestIdentity_ManageIdentitiesTrue() throws Exception {
    testCreateTestIdentity(new PrincipalKeyCredential("principal", "password"), Boolean.TRUE);
  }

  @Test
  public void testCreateTestIdentity_ManageIdentitiesFalse() throws Exception {
    testCreateTestIdentity(new PrincipalKeyCredential("principal", "password"), Boolean.FALSE);
  }

  @Test(expected = KerberosMissingAdminCredentialsException.class)
  public void testCreateTestIdentityNoCredentials_ManageIdentitiesDefault() throws Exception {
    testCreateTestIdentity(null, null);
  }

  @Test(expected = KerberosMissingAdminCredentialsException.class)
  public void testCreateTestIdentityNoCredentials_ManageIdentitiesTrue() throws Exception {
    testCreateTestIdentity(null, Boolean.TRUE);
  }

  @Test
  public void testCreateTestIdentityNoCredentials_ManageIdentitiesFalse() throws Exception {
    testCreateTestIdentity(null, Boolean.FALSE);
  }

  @Test
  public void testDeleteTestIdentity() throws Exception {
    testDeleteTestIdentity(new PrincipalKeyCredential("principal", "password"));
  }

  @Test(expected = IllegalArgumentException.class)
  public void testGetActiveIdentities_MissingCluster() throws Exception {
    testGetActiveIdentities(null, null, null, null, true, SecurityType.KERBEROS);
  }

  @Test
  public void testGetActiveIdentities_SecurityTypeKerberos_All() throws Exception {
    testGetActiveIdentities_All(SecurityType.KERBEROS);
  }

  @Test
  public void testGetActiveIdentities_SecurityTypeNone_All() throws Exception {
    testGetActiveIdentities_All(SecurityType.NONE);
  }

  @Test
  public void testGetActiveIdentities_SingleHost() throws Exception {
    Map<String, Collection<KerberosIdentityDescriptor>> identities = testGetActiveIdentities("c1", "host1", null, null, true, SecurityType.KERBEROS);

    Assert.assertNotNull(identities);
    Assert.assertEquals(1, identities.size());

    Collection<KerberosIdentityDescriptor> hostIdentities;

    hostIdentities = identities.get("host1");
    Assert.assertNotNull(hostIdentities);
    Assert.assertEquals(3, hostIdentities.size());

    validateIdentities(hostIdentities, new HashMap<String, Map<String, Object>>() {{
      put("identity1", new HashMap<String, Object>() {
        {
          put("principal_name", "service1/host1@EXAMPLE.COM");
          put("principal_type", KerberosPrincipalType.SERVICE);
          put("principal_configuration", "service1-site/component1.kerberos.principal");
          put("principal_local_username", "service1");
          put("keytab_file", "${keytab_dir}/service1.keytab");
          put("keytab_owner_name", "service1");
          put("keytab_owner_access", "rw");
          put("keytab_group_name", "hadoop");
          put("keytab_group_access", "");
          put("keytab_configuration", "service1-site/component1.keytab.file");
          put("keytab_cachable", false);
        }
      });

      put("identity2", new HashMap<String, Object>() {
        {
          put("principal_name", "component2/host1@EXAMPLE.COM");
          put("principal_type", KerberosPrincipalType.SERVICE);
          put("principal_configuration", "service2-site/component2.kerberos.principal");
          put("principal_local_username", "service2");
          put("keytab_file", "${keytab_dir}/service2.keytab");
          put("keytab_owner_name", "service2");
          put("keytab_owner_access", "rw");
          put("keytab_group_name", "hadoop");
          put("keytab_group_access", "");
          put("keytab_configuration", "service2-site/component2.keytab.file");
          put("keytab_cachable", false);
        }
      });

      put("identity3", new HashMap<String, Object>() {
        {
          put("principal_name", "service1/host1@EXAMPLE.COM");
          put("principal_type", KerberosPrincipalType.SERVICE);
          put("principal_configuration", "service1-site/service1.kerberos.principal");
          put("principal_local_username", "service1");
          put("keytab_file", "${keytab_dir}/service1.service.keytab");
          put("keytab_owner_name", "service1");
          put("keytab_owner_access", "rw");
          put("keytab_group_name", "hadoop");
          put("keytab_group_access", "");
          put("keytab_configuration", "service1-site/service1.keytab.file");
          put("keytab_cachable", false);
        }
      });
    }});
  }

  @Test
  public void addAmbariServerIdentity_CreateAmbariPrincipal() throws Exception {
    addAmbariServerIdentity(Collections.singletonMap(KerberosHelper.CREATE_AMBARI_PRINCIPAL, "true"));
  }

  @Test
  public void addAmbariServerIdentity_DoNotCreateAmbariPrincipal() throws Exception {
    addAmbariServerIdentity(Collections.singletonMap(KerberosHelper.CREATE_AMBARI_PRINCIPAL, "false"));
  }

  @Test
  public void addAmbariServerIdentity_MissingProperty() throws Exception {
    addAmbariServerIdentity(Collections.singletonMap("not_create_ambari_principal", "value"));
  }

  @Test
  public void addAmbariServerIdentity_MissingKerberosEnv() throws Exception {
    addAmbariServerIdentity(null);
  }

  @Test
  public void testGetActiveIdentities_SingleService() throws Exception {
    Map<String, Collection<KerberosIdentityDescriptor>> identities = testGetActiveIdentities("c1", null, "SERVICE1", null, true, SecurityType.KERBEROS);

    Assert.assertNotNull(identities);
    Assert.assertEquals(3, identities.size());

    Collection<KerberosIdentityDescriptor> hostIdentities;

    hostIdentities = identities.get("host1");
    Assert.assertNotNull(hostIdentities);
    Assert.assertEquals(2, hostIdentities.size());

    validateIdentities(hostIdentities, new HashMap<String, Map<String, Object>>() {{
      put("identity1", new HashMap<String, Object>() {
        {
          put("principal_name", "service1/host1@EXAMPLE.COM");
          put("principal_type", KerberosPrincipalType.SERVICE);
          put("principal_configuration", "service1-site/component1.kerberos.principal");
          put("principal_local_username", "service1");
          put("keytab_file", "${keytab_dir}/service1.keytab");
          put("keytab_owner_name", "service1");
          put("keytab_owner_access", "rw");
          put("keytab_group_name", "hadoop");
          put("keytab_group_access", "");
          put("keytab_configuration", "service1-site/component1.keytab.file");
          put("keytab_cachable", false);
        }
      });

      put("identity3", new HashMap<String, Object>() {
        {
          put("principal_name", "service1/host1@EXAMPLE.COM");
          put("principal_type", KerberosPrincipalType.SERVICE);
          put("principal_configuration", "service1-site/service1.kerberos.principal");
          put("principal_local_username", "service1");
          put("keytab_file", "${keytab_dir}/service1.service.keytab");
          put("keytab_owner_name", "service1");
          put("keytab_owner_access", "rw");
          put("keytab_group_name", "hadoop");
          put("keytab_group_access", "");
          put("keytab_configuration", "service1-site/service1.keytab.file");
          put("keytab_cachable", false);
        }
      });
    }});

