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
   * <p>
   * http://www.apache.org/licenses/LICENSE-2.0
   * <p>
   * 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.atlas.utils;
  
  import org.apache.atlas.ApplicationProperties;
  import org.apache.atlas.security.SecurityUtil;
  import org.apache.commons.configuration.Configuration;
  import org.apache.commons.lang.StringUtils;
  import org.apache.hadoop.security.UserGroupInformation;
  import org.apache.kafka.clients.admin.AdminClient;
  import org.apache.kafka.clients.admin.CreateTopicsResult;
  import org.apache.kafka.clients.admin.DescribeTopicsResult;
  import org.apache.kafka.clients.admin.ListTopicsOptions;
  import org.apache.kafka.clients.admin.ListTopicsResult;
  import org.apache.kafka.clients.admin.NewTopic;
  import org.apache.kafka.clients.admin.TopicDescription;
  import org.apache.kafka.common.KafkaFuture;
  import org.apache.kafka.common.TopicPartitionInfo;
  import org.apache.kafka.common.errors.TopicExistsException;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.Properties;
  import java.util.concurrent.ExecutionException;
  import java.util.stream.Collectors;
  
  import static org.apache.atlas.security.SecurityProperties.HADOOP_SECURITY_CREDENTIAL_PROVIDER_PATH;
  
  public class KafkaUtils implements AutoCloseable {
      private static final Logger LOG = LoggerFactory.getLogger(KafkaUtils.class);
  
      public static final String ATLAS_KAFKA_PROPERTY_PREFIX     = "atlas.kafka";
      public static final String KAFKA_SASL_JAAS_CONFIG_PROPERTY = "sasl.jaas.config";
      public static final String JAAS_PASSWORD_SUFFIX            = "password";
  
      private static final String JAAS_CONFIG_PREFIX_PARAM                     = "atlas.jaas";
      private static final String JAAS_CONFIG_LOGIN_MODULE_NAME_PARAM          = "loginModuleName";
      private static final String JAAS_CONFIG_LOGIN_MODULE_CONTROL_FLAG_PARAM  = "loginModuleControlFlag";
      private static final String JAAS_DEFAULT_LOGIN_MODULE_CONTROL_FLAG       = "required";
      private static final String JAAS_VALID_LOGIN_MODULE_CONTROL_FLAG_OPTIONS = "optional|requisite|sufficient|required";
      private static final String JAAS_CONFIG_LOGIN_OPTIONS_PREFIX             = "option";
      private static final String JAAS_PRINCIPAL_PROP                          = "principal";
      private static final String JAAS_DEFAULT_CLIENT_NAME                     = "KafkaClient";
      private static final String JAAS_TICKET_BASED_CLIENT_NAME                = "ticketBased-KafkaClient";
      private static final String IMPORT_INTERNAL_TOPICS                       = "atlas.kafka.bridge.enable.internal.topics.import";
      private static final String JAAS_MASK_PASSWORD                           = "********";
  
      protected final Properties  kafkaConfiguration;
      protected final AdminClient adminClient;
      protected final boolean     importInternalTopics;
  
      public KafkaUtils(Configuration atlasConfiguration) {
          LOG.debug("==> KafkaUtils() ");
  
          this.kafkaConfiguration = ApplicationProperties.getSubsetAsProperties(atlasConfiguration, ATLAS_KAFKA_PROPERTY_PREFIX);
  
          setKafkaJAASProperties(atlasConfiguration, kafkaConfiguration);
  
          adminClient          = AdminClient.create(this.kafkaConfiguration);
          importInternalTopics = atlasConfiguration.getBoolean(IMPORT_INTERNAL_TOPICS, false);
  
          LOG.debug("<== KafkaUtils() ");
      }
  
      public static void setKafkaJAASProperties(Configuration configuration, Properties kafkaProperties) {
          LOG.debug("==> KafkaUtils.setKafkaJAASProperties()");
  
          if (kafkaProperties.containsKey(KAFKA_SASL_JAAS_CONFIG_PROPERTY)) {
              LOG.debug("JAAS config is already set, returning");
  
              return;
          }
  
          Properties jaasConfig = ApplicationProperties.getSubsetAsProperties(configuration, JAAS_CONFIG_PREFIX_PARAM);
          // JAAS Configuration is present then update set those properties in sasl.jaas.config
  
          if (jaasConfig != null && !jaasConfig.isEmpty()) {
              String jaasClientName = JAAS_DEFAULT_CLIENT_NAME;
  
              // Required for backward compatability for Hive CLI
              if (!isLoginKeytabBased() && isLoginTicketBased()) {
                  LOG.debug("Checking if ticketBased-KafkaClient is set");
  
                  // if ticketBased-KafkaClient property is not specified then use the default client name
                  String        ticketBasedConfigPrefix = JAAS_CONFIG_PREFIX_PARAM + "." + JAAS_TICKET_BASED_CLIENT_NAME;
                  Configuration ticketBasedConfig       = configuration.subset(ticketBasedConfigPrefix);
  
                  if (ticketBasedConfig != null && !ticketBasedConfig.isEmpty()) {
                      LOG.debug("ticketBased-KafkaClient JAAS configuration is set, using it");
  
                      jaasClientName = JAAS_TICKET_BASED_CLIENT_NAME;
                  } else {
                      LOG.info("UserGroupInformation.isLoginTicketBased is true, but no JAAS configuration found for client {}. Will use JAAS configuration of client {}", JAAS_TICKET_BASED_CLIENT_NAME, jaasClientName);
                  }
              }
  
              String keyPrefix       = jaasClientName + ".";
              String keyParam        = keyPrefix + JAAS_CONFIG_LOGIN_MODULE_NAME_PARAM;
              String loginModuleName = jaasConfig.getProperty(keyParam);
  
              if (loginModuleName == null) {
                  LOG.error("Unable to add JAAS configuration for client [{}] as it is missing param [{}]. Skipping JAAS config for [{}]", jaasClientName, keyParam, jaasClientName);
  
                  return;
              }
  
              keyParam = keyPrefix + JAAS_CONFIG_LOGIN_MODULE_CONTROL_FLAG_PARAM;
  
              String controlFlag = jaasConfig.getProperty(keyParam);
  
              if (StringUtils.isEmpty(controlFlag)) {
                  String validValues = JAAS_VALID_LOGIN_MODULE_CONTROL_FLAG_OPTIONS;
  
                  controlFlag = JAAS_DEFAULT_LOGIN_MODULE_CONTROL_FLAG;
  
                  LOG.warn("Unknown JAAS configuration value for ({}) = [{}], valid value are [{}] using the default value, REQUIRED", keyParam, controlFlag, validValues);
              }
  
              String        optionPrefix       = keyPrefix + JAAS_CONFIG_LOGIN_OPTIONS_PREFIX + ".";
              String        principalOptionKey = optionPrefix + JAAS_PRINCIPAL_PROP;
              String        passwordOptionKey  = optionPrefix + JAAS_PASSWORD_SUFFIX;
              int           optionPrefixLen    = optionPrefix.length();
              StringBuilder optionStringBuffer = new StringBuilder();
  
              for (String key : jaasConfig.stringPropertyNames()) {
                  if (key.startsWith(optionPrefix)) {
                      String optionVal = jaasConfig.getProperty(key);
  
                      if (optionVal != null) {
                          optionVal = optionVal.trim();
  
                          try {
                              if (key.equalsIgnoreCase(principalOptionKey)) {
                                  optionVal = org.apache.hadoop.security.SecurityUtil.getServerPrincipal(optionVal, (String) null);
                              }
                          } catch (IOException e) {
                              LOG.warn("Failed to build serverPrincipal. Using provided value:[{}]", optionVal);
                          }
  
                          if (key.equalsIgnoreCase(passwordOptionKey)) {
                              String jaasKafkaClientConfigurationProperty = "atlas.jaas.KafkaClient.option.password";
  
                              if (JAAS_MASK_PASSWORD.equals(configuration.getString(jaasKafkaClientConfigurationProperty))) {
                                  try {
                                      optionVal = SecurityUtil.getPassword(configuration, jaasKafkaClientConfigurationProperty, HADOOP_SECURITY_CREDENTIAL_PROVIDER_PATH);
                                  } catch (Exception e) {
                                      LOG.error("Error in getting secure password ", e);
                                  }
                              }
                          }
  
                          optionVal = surroundWithQuotes(optionVal);
  
                          optionStringBuffer.append(String.format(" %s=%s", key.substring(optionPrefixLen), optionVal));
                      }
                  }
              }
  
              String newJaasProperty = String.format("%s %s %s ;", loginModuleName.trim(), controlFlag, optionStringBuffer);
  
              kafkaProperties.put(KAFKA_SASL_JAAS_CONFIG_PROPERTY, newJaasProperty);
          }
  
          LOG.debug("<== KafkaUtils.setKafkaJAASProperties()");
      }
  
      public static boolean isLoginTicketBased() {
          boolean ret = false;
  
          try {
              ret = UserGroupInformation.isLoginTicketBased();
          } catch (Exception excp) {
              LOG.warn("Error in determining ticket-cache for KafkaClient-JAAS config", excp);
          }
  
          return ret;
      }
  
      static boolean isLoginKeytabBased() {
          boolean ret = false;
  
          try {
              ret = UserGroupInformation.isLoginKeytabBased();
          } catch (Exception excp) {
              LOG.warn("Error in determining keytab for KafkaClient-JAAS config", excp);
          }
  
          return ret;
      }
  
      static String surroundWithQuotes(String optionVal) {
          if (StringUtils.isEmpty(optionVal)) {
              return optionVal;
          }
  
          // Enclose property values in double quotes, so that Kafka can parse it.
          // Escape all double quotes that may occur in the property value.
          String doubleQuoteEscaped = optionVal.replace("\"", "\\\"");
          return String.format("\"%s\"", doubleQuoteEscaped);
      }
  
      public void createTopics(List<String> topicNames, int numPartitions, int replicationFactor) throws TopicExistsException, ExecutionException, InterruptedException {
          LOG.debug("==> createTopics() ");
  
          List<NewTopic> newTopicList = topicNames.stream()
                  .map(topicName -> new NewTopic(topicName, numPartitions, (short) replicationFactor))
                  .collect(Collectors.toList());
  
          CreateTopicsResult             createTopicsResult = adminClient.createTopics(newTopicList);
          Map<String, KafkaFuture<Void>> futureMap          = createTopicsResult.values();
  
          for (Map.Entry<String, KafkaFuture<Void>> futureEntry : futureMap.entrySet()) {
              KafkaFuture<Void> future = futureEntry.getValue();
  
              future.get();
          }
  
          LOG.debug("<== createTopics() ");
      }
  
      public List<String> listAllTopics() throws ExecutionException, InterruptedException {
          LOG.debug("==> KafkaUtils.listAllTopics() ");
  
          ListTopicsResult listTopicsResult = adminClient.listTopics((new ListTopicsOptions()).listInternal(importInternalTopics));
          List<String>     ret              = new ArrayList<>(listTopicsResult.names().get());
  
          LOG.debug("<== KafkaUtils.listAllTopics() ");
  
          return ret;
      }
  
      public Integer getPartitionCount(String topicName) throws ExecutionException, InterruptedException {
          LOG.debug("==> KafkaUtils.getPartitionCount({})", topicName);
  
          Integer                  ret           = null;
          List<TopicPartitionInfo> partitionList = getPartitionList(topicName);
  
          if (partitionList != null) {
              ret = partitionList.size();
          }
  
          LOG.debug("<== KafkaUtils.getPartitionCount returning for topic {} with count {}", topicName, ret);
  
          return ret;
      }
  
      public Integer getReplicationFactor(String topicName) throws ExecutionException, InterruptedException {
          LOG.debug("==> KafkaUtils.getReplicationFactor({})", topicName);
  
          Integer                  ret           = null;
          List<TopicPartitionInfo> partitionList = getPartitionList(topicName);
  
          if (partitionList != null) {
              ret = partitionList.stream().mapToInt(x -> x.replicas().size()).max().orElse(1);
          }
  
          LOG.debug("<== KafkaUtils.getReplicationFactor returning for topic {} with replicationFactor {}", topicName, ret);
  
          return ret;
      }
  
      public void close() {
          LOG.debug("==> KafkaUtils.close()");
  
          if (adminClient != null) {
              adminClient.close();
          }
  
          LOG.debug("<== KafkaUtils.close()");
      }
  
      private List<TopicPartitionInfo> getPartitionList(String topicName) throws ExecutionException, InterruptedException {
          List<TopicPartitionInfo> ret                  = null;
          DescribeTopicsResult     describeTopicsResult = adminClient.describeTopics(Collections.singleton(topicName));
  
          if (describeTopicsResult != null) {
              Map<String, KafkaFuture<TopicDescription>> futureMap = describeTopicsResult.values();
  
              for (Map.Entry<String, KafkaFuture<TopicDescription>> futureEntry : futureMap.entrySet()) {
                  KafkaFuture<TopicDescription> topicDescriptionFuture = futureEntry.getValue();
                  TopicDescription              topicDescription       = topicDescriptionFuture.get();
  
                  ret = topicDescription.partitions();
              }
          }
  
          return ret;
      }
  }
  