content: /*
   * 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.events.publishers;
  
  import static org.apache.ambari.server.controller.KerberosHelperImpl.CHECK_KEYTABS;
  import static org.apache.ambari.server.controller.KerberosHelperImpl.REMOVE_KEYTAB;
  import static org.apache.ambari.server.controller.KerberosHelperImpl.SET_KEYTAB;
  
  import java.io.BufferedInputStream;
  import java.io.File;
  import java.io.FileInputStream;
  import java.io.IOException;
  import java.nio.file.Paths;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.TreeMap;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ExecutionException;
  
  import javax.annotation.Nullable;
  
  import org.apache.ambari.server.AmbariException;
  import org.apache.ambari.server.AmbariRuntimeException;
  import org.apache.ambari.server.agent.AgentCommand;
  import org.apache.ambari.server.agent.CancelCommand;
  import org.apache.ambari.server.agent.ExecutionCommand;
  import org.apache.ambari.server.agent.stomp.AgentConfigsHolder;
  import org.apache.ambari.server.agent.stomp.dto.ExecutionCommandsCluster;
  import org.apache.ambari.server.events.ExecutionCommandEvent;
  import org.apache.ambari.server.orm.dao.HostRoleCommandDAO;
  import org.apache.ambari.server.serveraction.kerberos.KerberosIdentityDataFileReader;
  import org.apache.ambari.server.serveraction.kerberos.KerberosServerAction;
  import org.apache.ambari.server.serveraction.kerberos.stageutils.KerberosKeytabController;
  import org.apache.ambari.server.serveraction.kerberos.stageutils.ResolvedKerberosKeytab;
  import org.apache.ambari.server.serveraction.kerberos.stageutils.ResolvedKerberosPrincipal;
  import org.apache.ambari.server.state.Clusters;
  import org.apache.ambari.server.state.DesiredConfig;
  import org.apache.ambari.server.state.kerberos.KerberosIdentityDescriptor;
  import org.apache.ambari.server.utils.StageUtils;
  import org.apache.ambari.server.utils.ThreadPools;
  import org.apache.commons.codec.binary.Base64;
  import org.apache.commons.codec.digest.DigestUtils;
  import org.apache.commons.collections4.CollectionUtils;
  import org.apache.commons.io.IOUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import com.google.common.collect.ArrayListMultimap;
  import com.google.common.collect.Multimap;
  import com.google.inject.Inject;
  import com.google.inject.Singleton;
  
  @Singleton
  public class AgentCommandsPublisher {
    private static final Logger LOG = LoggerFactory.getLogger(AgentCommandsPublisher.class);
  
    @Inject
    private KerberosKeytabController kerberosKeytabController;
  
    @Inject
    private Clusters clusters;
  
    @Inject
    private HostRoleCommandDAO hostRoleCommandDAO;
  
    @Inject
    private STOMPUpdatePublisher STOMPUpdatePublisher;
  
    @Inject
    private AgentConfigsHolder agentConfigsHolder;
  
    @Inject
    private ThreadPools threadPools;
  
  
    public void sendAgentCommand(Multimap<Long, AgentCommand> agentCommands) throws AmbariRuntimeException {
      if (agentCommands != null && !agentCommands.isEmpty()) {
        Map<Long, TreeMap<String, ExecutionCommandsCluster>> executionCommandsClusters = new ConcurrentHashMap<>();
        Map<Long, Map<String, DesiredConfig>> clusterDesiredConfigs = new ConcurrentHashMap<>();
  
        try {
          threadPools.getAgentPublisherCommandsPool().submit(() -> {
            agentCommands.entries().stream().parallel().forEach(acHostEntry -> {
              Long hostId = acHostEntry.getKey();
              AgentCommand ac = acHostEntry.getValue();
              Long clusterId = null;
  
              if (ac instanceof ExecutionCommand) {
                try {
                  clusterId = Long.valueOf(((ExecutionCommand)ac).getClusterId());
                  if (clusterId >= 0) {
                    if (!clusterDesiredConfigs.containsKey(clusterId)) {
                      clusterDesiredConfigs.put(clusterId, clusters.getCluster(clusterId).getDesiredConfigs());
                    }
                  } else {
                    LOG.warn("The cluster not found or has not been created yet. clusterID={}.", clusterId);
                  }
                } catch (NumberFormatException|AmbariException e) {
                  LOG.error("Exception on sendAgentCommand", e);
                }
              }
  
              Map<String, DesiredConfig> desiredConfigs = (clusterId != null && clusterDesiredConfigs.containsKey(clusterId))
                ? clusterDesiredConfigs.get(clusterId)
                : null;
              populateExecutionCommandsClusters(executionCommandsClusters, hostId, ac, desiredConfigs);
            });
          }).get();
        } catch (InterruptedException|ExecutionException e) {
          LOG.error("Exception on sendAgentCommand", e);
        }
  
        try {
          threadPools.getAgentPublisherCommandsPool().submit(() -> {
            executionCommandsClusters.entrySet().stream().parallel().forEach(entry -> {
              STOMPUpdatePublisher.publish(new ExecutionCommandEvent(
                entry.getKey(),
                agentConfigsHolder.initializeDataIfNeeded(entry.getKey(), true).getTimestamp(),
                entry.getValue()
              ));
            });
          }).get();
        } catch (InterruptedException|ExecutionException e) {
          LOG.error("Exception on sendAgentCommand", e);
        }
      }
    }
  
    public void sendAgentCommand(Long hostId, AgentCommand agentCommand) throws AmbariRuntimeException {
      Multimap<Long, AgentCommand> agentCommands = ArrayListMultimap.create();
      agentCommands.put(hostId, agentCommand);
      sendAgentCommand(agentCommands);
    }
  
    private void populateExecutionCommandsClusters(
      Map<Long, TreeMap<String, ExecutionCommandsCluster>> executionCommandsClusters,
      Long hostId, AgentCommand ac,
      @Nullable Map<String, DesiredConfig> desiredConfigs) throws AmbariRuntimeException {
  
      if (LOG.isDebugEnabled()) {
        try {
          LOG.debug("Sending command string = " + StageUtils.jaxbToString(ac));
        } catch (Exception e) {
          throw new AmbariRuntimeException("Could not get jaxb string for command", e);
        }
      }
  
      switch (ac.getCommandType()) {
        case BACKGROUND_EXECUTION_COMMAND:
        case EXECUTION_COMMAND: {
          ExecutionCommand ec = (ExecutionCommand) ac;
          LOG.info("AgentCommandsPublisher.sendCommands: sending ExecutionCommand for host {}, role {}, roleCommand {}, and command ID {}, task ID {}",
              ec.getHostname(), ec.getRole(), ec.getRoleCommand(), ec.getCommandId(), ec.getTaskId());
          Map<String, String> hlp = ec.getCommandParams();
          if (hlp != null) {
            String customCommand = hlp.get("custom_command");
            if (SET_KEYTAB.equalsIgnoreCase(customCommand) || REMOVE_KEYTAB.equalsIgnoreCase(customCommand) || CHECK_KEYTABS.equalsIgnoreCase(customCommand)) {
              LOG.info(String.format("%s called", customCommand));
              try {
                injectKeytab(ec, customCommand, clusters.getHostById(hostId).getHostName(), desiredConfigs);
              } catch (IOException e) {
                throw new AmbariRuntimeException("Could not inject keytab into command", e);
              }
            }
          }
          String clusterName = ec.getClusterName();
          String clusterId = "-1";
          if (clusterName != null) {
            try {
              clusterId = Long.toString(clusters.getCluster(clusterName).getClusterId());
            } catch (AmbariException e) {
              throw new AmbariRuntimeException(e);
            }
          }
          ec.setClusterId(clusterId);
          prepareExecutionCommandsClusters(executionCommandsClusters, hostId, clusterId);
          executionCommandsClusters.get(hostId).get(clusterId).getExecutionCommands().add((ExecutionCommand) ac);
          break;
        }
        case CANCEL_COMMAND: {
          CancelCommand cc = (CancelCommand) ac;
          String clusterId = Long.toString(hostRoleCommandDAO.findByPK(cc.getTargetTaskId()).getStage().getClusterId());
          prepareExecutionCommandsClusters(executionCommandsClusters, hostId, clusterId);
          executionCommandsClusters.get(hostId).get(clusterId).getCancelCommands().add(cc);
          break;
        }
        default:
          LOG.error("There is no action for agent command ="
              + ac.getCommandType().name());
      }
    }
  
    private void prepareExecutionCommandsClusters(Map<Long, TreeMap<String, ExecutionCommandsCluster>> executionCommandsClusters,
                                                  Long hostId, String clusterId) {
      if (!executionCommandsClusters.containsKey(hostId)) {
        executionCommandsClusters.put(hostId, new TreeMap<>());
      }
      if (!executionCommandsClusters.get(hostId).containsKey(clusterId)) {
        executionCommandsClusters.get(hostId).put(clusterId, new ExecutionCommandsCluster(new ArrayList<>(),
            new ArrayList<>()));
      }
    }
  
    /**
     * Insert Kerberos keytab details into the ExecutionCommand for the SET_KEYTAB custom command if
     * any keytab details and associated data exists for the target host.
     *
     * @param ec the ExecutionCommand to update
     * @param command a name of the relevant keytab command
     * @param targetHost a name of the host the relevant command is destined for
     * @param desiredConfigs desired config map
     * @throws AmbariException
     */
    private void injectKeytab(ExecutionCommand ec, String command, String targetHost,
                              @Nullable Map<String, DesiredConfig> desiredConfigs) throws AmbariException {
      KerberosCommandParameterProcessor processor = KerberosCommandParameterProcessor.getInstance(command, clusters, ec, kerberosKeytabController);
      if (processor != null) {
        ec.setKerberosCommandParams(processor.process(targetHost, desiredConfigs));
      }
    }
  
    /**
     * KerberosCommandParameterProcessor is an abstract class providing common implementations for processing
     * the Kerberos command parameters.
     *
     * The Kerberos command parameters are processed differently depending on the operation
     * - set, check, or remove keytab files.
     */
    private static abstract class KerberosCommandParameterProcessor {
      protected final Clusters clusters;
  
      protected final ExecutionCommand executionCommand;
  
      protected final KerberosKeytabController kerberosKeytabController;
  
      protected List<Map<String, String>> kcp;
  
      protected KerberosCommandParameterProcessor(Clusters clusters, ExecutionCommand executionCommand, KerberosKeytabController kerberosKeytabController) {
        this.clusters = clusters;
        this.executionCommand = executionCommand;
        this.kerberosKeytabController = kerberosKeytabController;
        kcp = executionCommand.getKerberosCommandParams();
      }
  
      /**
       * Factory method to return the appropriate KerberosCommandParameterProcessor instance based
       * on the command being executed.
       *
       * @param command the command being executed
       * @param clusters the clusters helper class
       * @param executionCommand the execution command structure
       * @param kerberosKeytabController the keytab controller helper class
       * @return
       */
      public static KerberosCommandParameterProcessor getInstance(String command, Clusters clusters, ExecutionCommand executionCommand, KerberosKeytabController kerberosKeytabController) {
        if (SET_KEYTAB.equalsIgnoreCase(command)) {
          return new SetKeytabCommandParameterProcessor(clusters, executionCommand, kerberosKeytabController);
        }
        if (CHECK_KEYTABS.equalsIgnoreCase(command)) {
          return new CheckKeytabsCommandParameterProcessor(clusters, executionCommand, kerberosKeytabController);
        }
  
        if (REMOVE_KEYTAB.equalsIgnoreCase(command)) {
          return new RemoveKeytabCommandParameterProcessor(clusters, executionCommand, kerberosKeytabController);
        }
  
        return null;
      }
  
      /**
       * Performs the default behavior for processing the relevant Kerberos identities and generating the
       * Kerberos-specific command details to send to the agent.
       *
       * @param targetHost the hostname of the target host
       * @param desiredConfigs desired config map
       * @return a map of propoperties to set as the Kerberos command parameters
       * @throws AmbariException
       */
      public List<Map<String, String>> process(String targetHost,
                                               @Nullable Map<String, DesiredConfig> desiredConfigs) throws AmbariException {
        KerberosServerAction.KerberosCommandParameters kerberosCommandParameters = new KerberosServerAction.KerberosCommandParameters(executionCommand);
        Map<String, ? extends Collection<String>> serviceComponentFilter = getServiceComponentFilter(kerberosCommandParameters.getServiceComponentFilter());
        final Collection<KerberosIdentityDescriptor> serviceIdentities = serviceComponentFilter == null
          ? null
          : kerberosKeytabController.getServiceIdentities(executionCommand.getClusterName(), serviceComponentFilter.keySet(), desiredConfigs);
        final Set<ResolvedKerberosKeytab> keytabsToInject = kerberosKeytabController.getFilteredKeytabs(serviceIdentities, kerberosCommandParameters.getHostFilter(), kerberosCommandParameters.getIdentityFilter());
  
        keytabsToInject.forEach(resolvedKeytab -> resolvedKeytab.getPrincipals().forEach(
          resolvedPrincipal -> {
            String hostName = resolvedPrincipal.getHostName();
            if (targetHost.equalsIgnoreCase(hostName)) {
              try {
                process(targetHost, resolvedKeytab, resolvedPrincipal, serviceComponentFilter);
              } catch (IOException e) {
                throw new AmbariRuntimeException("Could not inject keytabs to enable kerberos", e);
              }
            }
          }
        ));
  
        return kcp;
      }
  
      /**
       * Performs the default behavior for processing the details of a particular Kerberos identity to
       * be added to the Kerberos command parameters.
       *
       * Implementations will override this method to perform specified tasks.
       *
       * @param hostName               the target hostname
       * @param resolvedKeytab         the relevant keytab file details
       * @param resolvedPrincipal      the relevant principal details
       * @param serviceComponentFilter the filter used to determine if the current Kerberos identity
       *                               should be processed
       * @throws IOException
       */
      protected void process(String hostName, ResolvedKerberosKeytab resolvedKeytab,
                             ResolvedKerberosPrincipal resolvedPrincipal,
                             Map<String, ? extends Collection<String>> serviceComponentFilter) throws IOException {
        Map<String, String> keytabMap = new HashMap<>();
        keytabMap.put(KerberosIdentityDataFileReader.HOSTNAME, hostName);
        keytabMap.put(KerberosIdentityDataFileReader.PRINCIPAL, resolvedPrincipal.getPrincipal());
        keytabMap.put(KerberosIdentityDataFileReader.KEYTAB_FILE_PATH, resolvedKeytab.getFile());
        kcp.add(keytabMap);
      }
  
      /**
       * Given a service/component filter, processes it as needed.
       * <p>
       * See overridden methods for more details.
       *
       * @param serviceComponentFilter a map of service to components indicate the services and
       *                               components to include in an operation
       * @return a map of service to components indicate the services and components to include in
       * the operation
       * @throws AmbariException
       */
      protected Map<String, ? extends Collection<String>> getServiceComponentFilter(Map<String, ? extends Collection<String>> serviceComponentFilter) throws AmbariException {
        return serviceComponentFilter;
      }
    }
  
    /**
     * SetKeytabCommandParameterProcessor is an implementation of {@link KerberosCommandParameterProcessor}
     * that handles the case for setting keytab files.
     * <p>
     * Specifically, this implementation add addition the keytab file details and its contents to the
     * command parameters. It also only performs operations only for services and components that are
     * known to be installed; therefore, the service/component filter may be altered to enforce this.
     */
    private static class SetKeytabCommandParameterProcessor extends KerberosCommandParameterProcessor {
  
      private final String dataDir;
  
      private SetKeytabCommandParameterProcessor(Clusters clusters, ExecutionCommand executionCommand, KerberosKeytabController kerberosKeytabController) {
        super(clusters, executionCommand, kerberosKeytabController);
        dataDir = executionCommand.getCommandParams().get(KerberosServerAction.DATA_DIRECTORY);
      }
  
      @Override
      protected void process(String hostName, ResolvedKerberosKeytab resolvedKeytab, ResolvedKerberosPrincipal resolvedPrincipal, Map<String, ? extends Collection<String>> serviceComponentFilter) throws IOException {
        if (dataDir != null) {
          String principal = resolvedPrincipal.getPrincipal();
          String keytabFilePath = resolvedKeytab.getFile();
          LOG.info("Processing principal {} for host {} and keytab file path {}", principal, hostName, keytabFilePath);
  
          if (keytabFilePath != null) {
            String sha1Keytab = DigestUtils.sha256Hex(keytabFilePath);
            File keytabFile = Paths.get(dataDir, hostName, sha1Keytab).toFile();
  
            if (keytabFile.canRead()) {
              Map<String, String> keytabMap = new HashMap<>();
  
              keytabMap.put(KerberosIdentityDataFileReader.HOSTNAME, hostName);
              keytabMap.put(KerberosIdentityDataFileReader.PRINCIPAL, principal);
              keytabMap.put(KerberosIdentityDataFileReader.KEYTAB_FILE_PATH, keytabFilePath);
              keytabMap.put(KerberosIdentityDataFileReader.KEYTAB_FILE_OWNER_NAME, resolvedKeytab.getOwnerName());
              keytabMap.put(KerberosIdentityDataFileReader.KEYTAB_FILE_OWNER_ACCESS, resolvedKeytab.getOwnerAccess());
              keytabMap.put(KerberosIdentityDataFileReader.KEYTAB_FILE_GROUP_NAME, resolvedKeytab.getGroupName());
              keytabMap.put(KerberosIdentityDataFileReader.KEYTAB_FILE_GROUP_ACCESS, resolvedKeytab.getGroupAccess());
  
              BufferedInputStream bufferedIn = new BufferedInputStream(new FileInputStream(keytabFile));
              byte[] keytabContent;
              try {
                keytabContent = IOUtils.toByteArray(bufferedIn);
              } finally {
                bufferedIn.close();
              }
              String keytabContentBase64 = Base64.encodeBase64String(keytabContent);
              keytabMap.put(KerberosServerAction.KEYTAB_CONTENT_BASE64, keytabContentBase64);
  
              kcp.add(keytabMap);
            } else {
              LOG.warn("Keytab file for principal {} and host {} can not to be read at path {}",
                  principal, hostName, keytabFile.getAbsolutePath());
            }
          }
        }
      }
  
      @Override
      protected Map<String, ? extends Collection<String>> getServiceComponentFilter(Map<String, ? extends Collection<String>> serviceComponentFilter)
          throws AmbariException {
        return kerberosKeytabController.adjustServiceComponentFilter(clusters.getCluster(executionCommand.getClusterName()), false, serviceComponentFilter);
      }
    }
  
    /**
     * CheckKeytabsCommandParameterProcessor is an implementation of {@link KerberosCommandParameterProcessor}
     * that handles the case for checking the keytab files on the hosts of the cluster.
     */
    private static class CheckKeytabsCommandParameterProcessor extends KerberosCommandParameterProcessor {
  
      private CheckKeytabsCommandParameterProcessor(Clusters clusters, ExecutionCommand executionCommand, KerberosKeytabController kerberosKeytabController) {
        super(clusters, executionCommand, kerberosKeytabController);
      }
    }
  
    /**
     * RemoveKeytabCommandParameterProcessor is an implementation of {@link KerberosCommandParameterProcessor}
     * that handles the case for setting keytab files.
     * <p>
     * Specifically, performs operations any services and components; however only keytab files found
     * to no longer be needed are specified for removal.
     */
    private static class RemoveKeytabCommandParameterProcessor extends KerberosCommandParameterProcessor {
  
      private RemoveKeytabCommandParameterProcessor(Clusters clusters, ExecutionCommand executionCommand, KerberosKeytabController kerberosKeytabController) {
        super(clusters, executionCommand, kerberosKeytabController);
      }
  
      @Override
      protected void process(String hostName, ResolvedKerberosKeytab resolvedKeytab, ResolvedKerberosPrincipal resolvedPrincipal, Map<String, ? extends Collection<String>> serviceComponentFilter) throws IOException {
        if (shouldRemove(hostName, resolvedKeytab, resolvedPrincipal, serviceComponentFilter)) {
          super.process(hostName, resolvedKeytab, resolvedPrincipal, serviceComponentFilter);
        }
      }
  
      /**
       * Determines if the keytab file for a given Kerberos identitiy should be removed from the target
       * host.
       * <p>
       * This is determined by comparing the service/component filter with the metadata about the relavent
       * Kerberos identity. If it is determined that more components than the ones specified in the filer
       * are linked to the identity, than the keytab file will not be flagged for removal.
       *
       * @param hostname               the target hostname
       * @param resolvedKerberosKeytab the relevant keytab file details
       * @param resolvedPrincipal      the relevant principal details
       * @param serviceComponentFilter the filter used to determine if the current Kerberos identity
       *                               should be processed
       * @return <code>true</code>, if this keytab file should be removed; <code>false</code>, otherwise
       */
      private boolean shouldRemove(String hostname,
                                   ResolvedKerberosKeytab resolvedKerberosKeytab,
                                   ResolvedKerberosPrincipal resolvedPrincipal,
                                   Map<String, ? extends Collection<String>> serviceComponentFilter) {
        ResolvedKerberosKeytab existingResolvedKeytab = kerberosKeytabController.getKeytabByFile(resolvedKerberosKeytab.getFile());
  
        if (existingResolvedKeytab == null) {
          return true;
        }
  
        Set<ResolvedKerberosPrincipal> principals = existingResolvedKeytab.getPrincipals();
        for (ResolvedKerberosPrincipal principal : principals) {
          if (hostname.equals(principal.getHostName()) && principal.getPrincipal().equals(resolvedPrincipal.getPrincipal())) {
            Multimap<String, String> temp = principal.getServiceMapping();
  
            // Make a local copy so we do not edit the stored copy, since we do not know how it is stored...
            Map<String, Collection<String>> serviceMapping = (temp == null) ? new HashMap<>() : new HashMap<>(temp.asMap());
  
            // Prune off the services in the filter, or all if the filter it none.  If there are no
            // service mappings left, this keytab file can be removed...
            if (serviceComponentFilter == null) {
              serviceMapping.clear();
            } else {
              for (Map.Entry<String, ? extends Collection<String>> entry : serviceComponentFilter.entrySet()) {
                String service = entry.getKey();
                Collection<String> components = entry.getValue();
  
                if (serviceMapping.containsKey(service)) {
  
                  if (CollectionUtils.isEmpty(components) || CollectionUtils.isEmpty(serviceMapping.get(service))) {
                    // Remove all entries for the service...
                    serviceMapping.remove(service);
                  } else {
                    Collection<String> leftOver = new HashSet<String>(serviceMapping.get(service));
                    leftOver.removeAll(components);
  
                    if (CollectionUtils.isEmpty(leftOver)) {
                      serviceMapping.remove(service);
                    } else {
                      serviceMapping.put(service, leftOver);
                    }
                  }
                }
              }
            }
  
            // There are still service mappings for this keytab files, we cannot remove it.
            if (serviceMapping.size() > 0) {
              return false;
            }
          }
        }
  
        return true;
      }
    }
  }
  