content: package io.aiven.klaw.clusterapi.utils;
  
  import static io.aiven.klaw.clusterapi.models.enums.KafkaSupportedProtocol.PLAINTEXT;
  import static io.aiven.klaw.clusterapi.models.enums.KafkaSupportedProtocol.SSL;
  
  import com.google.common.base.Strings;
  import io.aiven.klaw.clusterapi.config.SslContextConfig;
  import io.aiven.klaw.clusterapi.models.enums.KafkaClustersType;
  import io.aiven.klaw.clusterapi.models.enums.KafkaSupportedProtocol;
  import java.security.MessageDigest;
  import java.security.NoSuchAlgorithmException;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Properties;
  import java.util.concurrent.TimeUnit;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.commons.lang3.tuple.Pair;
  import org.apache.kafka.clients.CommonClientConfigs;
  import org.apache.kafka.clients.admin.AdminClient;
  import org.apache.kafka.clients.admin.AdminClientConfig;
  import org.apache.kafka.common.config.SaslConfigs;
  import org.apache.kafka.common.config.SslConfigs;
  import org.apache.tomcat.util.codec.binary.Base64;
  import org.springframework.beans.factory.annotation.Autowired;
  import org.springframework.core.env.Environment;
  import org.springframework.http.HttpHeaders;
  import org.springframework.stereotype.Service;
  import org.springframework.web.client.RestTemplate;
  
  @Service
  @Slf4j
  public class ClusterApiUtils {
  
    private static final String PLAIN_REST_TEMPLATE = "PlainRestTemplate";
    private static final String SSL_REST_TEMPLATE = "SSLRestTemplate";
    public static final String HTTPS_PREFIX = "https://";
    public static final String HTTP_PREFIX = "http://";
    public static final String SHA_256 = "SHA_256";
    public static final String SHA_512 = "SHA_512";
    public static final String KAFKA_CONNECT_CREDENTIALS_PROPERTY_SFX =
        ".klaw.kafkaconnect.credentials";
  
    public static final String KAFKA_SR_CREDENTIALS_PROPERTY_SFX = ".klaw.schemaregistry.credentials";
    public static final String KAFKA_CONFLUENT_CLOUD_CREDENTIALS_PROPERTY_SFX =
        ".klaw.confluentcloud.credentials";
  
    private static MessageDigest messageDigest;
  
    static {
      try {
        messageDigest = MessageDigest.getInstance("SHA-256");
      } catch (NoSuchAlgorithmException e) {
        log.error("Error ", e);
      }
    }
  
    private final Environment env;
    private final Map<String, AdminClient> adminClientsMap;
  
    private final Map<String, RestTemplate> restTemplateMap;
  
    private final AdminClientProperties adminClientProperties;
  
    @Autowired private SslContextConfig sslContextConfig;
  
    @Autowired
    public ClusterApiUtils(Environment env, AdminClientProperties adminClientProperties) {
      this(env, adminClientProperties, new HashMap<>(), new HashMap<>());
    }
  
    public AdminClientProperties getAdminClientProperties() {
      return adminClientProperties;
    }
  
    ClusterApiUtils(
        Environment env,
        AdminClientProperties adminClientProperties,
        Map<String, AdminClient> adminClientsMap,
        Map<String, RestTemplate> restTemplateMap) {
      this.env = env;
      this.adminClientsMap = adminClientsMap;
      this.adminClientProperties = adminClientProperties;
      this.restTemplateMap = restTemplateMap;
    }
  
    //    public void removeSSLElementFromAdminClientMap(String protocol, String clusterName){
    //        log.info("Into removeSSLElementFromAdminClientMap");
    //        String adminClientKeyReq = protocol + clusterName;
    //        List<String> sslKeys = adminClientsMap.keySet().stream()
    //                .filter(adminClientKey -> adminClientKey.equals(adminClientKeyReq))
    //                .collect(Collectors.toList());
    //        sslKeys.forEach(adminClientsMap::remove);
    //    }
  
    public String getHash(String envHost) {
      return new String(Base64.encodeBase64(messageDigest.digest(envHost.getBytes()), false));
    }
  
    public AdminClient getAdminClient(
        String envHost, KafkaSupportedProtocol protocol, String clusterIdentification)
        throws Exception {
      log.info(
          "Host : {} Protocol {} clusterIdentification {}", envHost, protocol, clusterIdentification);
  
      AdminClient adminClient = null;
      String adminClientKey = protocol + clusterIdentification + getHash(envHost);
  
      try {
        switch (protocol) {
          case PLAINTEXT:
            if (!adminClientsMap.containsKey(adminClientKey)) {
              adminClient = AdminClient.create(getPlainProperties(envHost));
            } else {
              adminClient = adminClientsMap.get(adminClientKey);
            }
            break;
  
          case SSL:
            if (!adminClientsMap.containsKey(adminClientKey)) {
              adminClient = AdminClient.create(getSslProperties(envHost, clusterIdentification));
            } else {
              adminClient = adminClientsMap.get(adminClientKey);
            }
            break;
  
          case SASL_PLAIN:
            if (!adminClientsMap.containsKey(adminClientKey)) {
              adminClient =
                  AdminClient.create(getSaslPlainProperties(envHost, clusterIdentification));
            } else {
              adminClient = adminClientsMap.get(adminClientKey);
            }
            break;
  
          case SASL_SSL_PLAIN_MECHANISM:
            if (!adminClientsMap.containsKey(adminClientKey)) {
              adminClient =
                  AdminClient.create(
                      getSaslSsl_PlainMechanismProperties(envHost, clusterIdentification));
            } else {
              adminClient = adminClientsMap.get(adminClientKey);
            }
  
            break;
  
          case SASL_SSL_SCRAM_MECHANISM_256:
            if (!adminClientsMap.containsKey(adminClientKey)) {
              adminClient =
                  AdminClient.create(
                      getSaslSsl_ScramMechanismProperties(envHost, clusterIdentification, SHA_256));
            } else {
              adminClient = adminClientsMap.get(adminClientKey);
            }
  
            break;
  
          case SASL_SSL_SCRAM_MECHANISM_512:
            if (!adminClientsMap.containsKey(adminClientKey)) {
              adminClient =
                  AdminClient.create(
                      getSaslSsl_ScramMechanismProperties(envHost, clusterIdentification, SHA_512));
            } else {
              adminClient = adminClientsMap.get(adminClientKey);
            }
  
            break;
  
          case SASL_SSL_GSSAPI_MECHANISM:
            if (!adminClientsMap.containsKey(adminClientKey)) {
              adminClient =
                  AdminClient.create(
                      getSaslSsl_GSSAPIMechanismProperties(envHost, clusterIdentification));
            } else {
              adminClient = adminClientsMap.get(adminClientKey);
            }
            break;
  
          case SASL_PLAIN_OAUTHBEARER:
            if (!adminClientsMap.containsKey(adminClientKey)) {
              adminClient =
                  AdminClient.create(
                      getSasl_OauthBearerMechanismProperties(envHost, clusterIdentification, false));
            } else {
              adminClient = adminClientsMap.get(adminClientKey);
            }
            break;
  
          case SASL_SSL_OAUTHBEARER:
            if (!adminClientsMap.containsKey(adminClientKey)) {
              adminClient =
                  AdminClient.create(
                      getSasl_OauthBearerMechanismProperties(envHost, clusterIdentification, true));
            } else {
              adminClient = adminClientsMap.get(adminClientKey);
            }
            break;
        }
      } catch (Exception exception) {
        log.error("Unable to create Admin client ", exception);
        exception.printStackTrace();
        throw new Exception("Cannot connect to cluster. Please contact Administrator.");
      }
  
      if (adminClient == null) {
        log.error("Cannot create Admin Client {} {}", envHost, protocol);
        throw new Exception("Cannot connect to cluster. Please contact Administrator.");
      }
  
      try {
        adminClient
            .listTopics()
            .names()
            .get(adminClientProperties.getTopicsTimeoutSecs(), TimeUnit.SECONDS);
        if (!adminClientsMap.containsKey(adminClientKey)) {
          adminClientsMap.put(adminClientKey, adminClient);
        }
        return adminClient;
      } catch (Exception e) {
        adminClientsMap.remove(adminClientKey);
        adminClient.close();
        log.error("Cannot create Admin Client {} {} {}", envHost, protocol, clusterIdentification, e);
        throw new Exception("Cannot connect to cluster. Please contact Administrator.");
      }
    }
  
    public Properties getPlainProperties(String environment) {
      Properties props = new Properties();
  
      props.put("bootstrap.servers", environment);
      setOtherConfig(props);
  
      return props;
    }
  
    public Properties getSslProperties(String environment, String clusterIdentification) {
      Properties props = getSslConfig(clusterIdentification);
  
      props.put("bootstrap.servers", environment);
      props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SSL");
      props.put(AdminClientConfig.CLIENT_ID_CONFIG, "klawclientssl");
      setOtherConfig(props);
  
      return props;
    }
  
    public Properties getSaslPlainProperties(String environment, String clusterIdentification) {
      Properties props = new Properties();
  
      props.put("bootstrap.servers", environment);
  
      try {
        props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SASL_PLAINTEXT");
        props.put(AdminClientConfig.CLIENT_ID_CONFIG, "klawclientsaslplain");
        setOtherConfig(props);
  
        if (!Strings.isNullOrEmpty(env.getProperty("kafkasasl.saslmechanism.plain"))) {
          props.put(SaslConfigs.SASL_MECHANISM, env.getProperty("kafkasasl.saslmechanism.plain"));
        }
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.plain"))) {
          props.put(
              SaslConfigs.SASL_JAAS_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.plain"));
        }
  
      } catch (Exception exception) {
        log.error("Error : Cannot set SASL PLAIN Config properties.", exception);
      }
  
      return props;
    }
  
    public Properties getSaslSsl_PlainMechanismProperties(
        String environment, String clusterIdentification) {
      Properties props = getSslConfig(clusterIdentification);
  
      try {
        props.put("bootstrap.servers", environment);
        props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SASL_SSL");
        props.put(AdminClientConfig.CLIENT_ID_CONFIG, "klawclientsaslsslplain");
        setOtherConfig(props);
  
        if (!Strings.isNullOrEmpty(env.getProperty("kafkasasl.saslmechanism.plain"))) {
          props.put(SaslConfigs.SASL_MECHANISM, env.getProperty("kafkasasl.saslmechanism.plain"));
        }
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.plain"))) {
          props.put(
              SaslConfigs.SASL_JAAS_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.plain"));
        }
      } catch (Exception exception) {
        log.error("Error : Cannot set SASL SSL PLAIN Config properties.", exception);
      }
  
      return props;
    }
  
    public Properties getSaslSsl_ScramMechanismProperties(
        String environment, String clusterIdentification, String algorithm) {
      Properties props = getSslConfig(clusterIdentification);
  
      try {
        props.put("bootstrap.servers", environment);
        props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SASL_SSL");
        props.put(AdminClientConfig.CLIENT_ID_CONFIG, "klawclientsaslsslscram");
        setOtherConfig(props);
  
        if (SHA_256.equals(algorithm)) {
          if (!Strings.isNullOrEmpty(env.getProperty("kafkasasl.saslmechanism.scram.256"))) {
            props.put(
                SaslConfigs.SASL_MECHANISM, env.getProperty("kafkasasl.saslmechanism.scram.256"));
          }
        } else if (SHA_512.equals(algorithm)) {
          if (!Strings.isNullOrEmpty(env.getProperty("kafkasasl.saslmechanism.scram.512"))) {
            props.put(
                SaslConfigs.SASL_MECHANISM, env.getProperty("kafkasasl.saslmechanism.scram.512"));
          }
        }
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.scram"))) {
          props.put(
              SaslConfigs.SASL_JAAS_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.scram"));
        }
      } catch (Exception exception) {
        log.error("Error : Cannot set SASL SSL SCRAM Config properties.", exception);
      }
  
      return props;
    }
  
    public Properties getSaslSsl_GSSAPIMechanismProperties(
        String environment, String clusterIdentification) {
      Properties props = getSslConfig(clusterIdentification);
  
      try {
        props.put("bootstrap.servers", environment);
        props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SASL_SSL");
        props.put(AdminClientConfig.CLIENT_ID_CONFIG, "klawclientsaslsslgssapi");
        setOtherConfig(props);
  
        if (!Strings.isNullOrEmpty(env.getProperty("kafkasasl.saslmechanism.gssapi"))) {
          props.put(SaslConfigs.SASL_MECHANISM, env.getProperty("kafkasasl.saslmechanism.gssapi"));
        }
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.gssapi"))) {
          props.put(
              SaslConfigs.SASL_JAAS_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.gssapi"));
        }
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(
                clusterIdentification.toLowerCase()
                    + ".kafkasasl.saslmechanism.gssapi.servicename"))) {
          props.put(
              SaslConfigs.SASL_KERBEROS_SERVICE_NAME,
              env.getProperty(
                  clusterIdentification.toLowerCase()
                      + ".kafkasasl.saslmechanism.gssapi.servicename"));
        }
      } catch (Exception exception) {
        log.error("Error : Cannot set SASL SSL GSSAPI Config properties.", exception);
      }
  
      return props;
    }
  
    public Properties getSasl_OauthBearerMechanismProperties(
        String environment, String clusterIdentification, boolean isSsl) {
      Properties props = new Properties();
  
      try {
        if (isSsl) {
          props = getSslConfig(clusterIdentification);
          props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SASL_SSL");
        } else {
          props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SASL_PLAINTEXT");
        }
        props.put("bootstrap.servers", environment);
  
        props.put(AdminClientConfig.CLIENT_ID_CONFIG, "klawclientsaslplainoauthbearer");
        setOtherConfig(props);
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(
                clusterIdentification.toLowerCase() + ".kafkasasl.saslmechanism.oauthbearer"))) {
          props.put(
              SaslConfigs.SASL_MECHANISM,
              env.getProperty(
                  clusterIdentification.toLowerCase() + ".kafkasasl.saslmechanism.oauthbearer"));
        }
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(
                clusterIdentification.toLowerCase()
                    + ".kafkasasl.sasl.oauthbearer.token.endpoint.url"))) {
          props.put(
              SaslConfigs.SASL_OAUTHBEARER_TOKEN_ENDPOINT_URL,
              env.getProperty(
                  clusterIdentification.toLowerCase()
                      + ".kafkasasl.sasl.oauthbearer.token.endpoint.url"));
        }
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(
                clusterIdentification.toLowerCase()
                    + ".kafkasasl.sasl.client.callback.handler.class"))) {
          props.put(
              SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS,
              env.getProperty(
                  clusterIdentification.toLowerCase()
                      + ".kafkasasl.sasl.client.callback.handler.class"));
        }
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(
                clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.oauthbearer"))) {
          props.put(
              SaslConfigs.SASL_JAAS_CONFIG,
              env.getProperty(
                  clusterIdentification.toLowerCase() + ".kafkasasl.jaasconfig.oauthbearer"));
        }
  
      } catch (Exception exception) {
        log.error("Error : Cannot set SASL SSL OAUTHBEARER Config properties.", exception);
      }
  
      return props;
    }
  
    public Properties getSslConfig(String clusterIdentification) {
      Properties props = new Properties();
  
      try {
        if (!Strings.isNullOrEmpty(
            env.getProperty(
                clusterIdentification.toLowerCase() + ".kafkassl.keystore.certificate.chain"))) {
          props.put(
              SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG,
              env.getProperty(
                  clusterIdentification.toLowerCase() + ".kafkassl.keystore.certificate.chain"));
        }
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.keystore.key"))) {
          props.put(
              SslConfigs.SSL_KEYSTORE_KEY_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.keystore.key"));
        }
        if (!Strings.isNullOrEmpty(
            env.getProperty(
                clusterIdentification.toLowerCase() + ".kafkassl.truststore.certificates"))) {
          props.put(
              SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG,
              env.getProperty(
                  clusterIdentification.toLowerCase() + ".kafkassl.truststore.certificates"));
        }
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.keystore.location"))) {
          props.put(
              SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.keystore.location"));
        }
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.keystore.pwd"))) {
          props.put(
              SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.keystore.pwd"));
        }
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.key.pwd"))) {
          props.put(
              SslConfigs.SSL_KEY_PASSWORD_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.key.pwd"));
        }
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.keystore.type"))) {
          props.put(
              SslConfigs.SSL_KEYSTORE_TYPE_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.keystore.type"));
        } else {
          props.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, "JKS");
        }
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.truststore.type"))) {
          props.put(
              SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.truststore.type"));
        } else {
          props.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, "JKS");
        }
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.truststore.location"))) {
          props.put(
              SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.truststore.location"));
        }
  
        if (!Strings.isNullOrEmpty(
            env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.truststore.pwd"))) {
          props.put(
              SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG,
              env.getProperty(clusterIdentification.toLowerCase() + ".kafkassl.truststore.pwd"));
        }
  
        props.put("ssl.enabled.protocols", "TLSv1.2,TLSv1.1");
        props.put("ssl.endpoint.identification.algorithm", "");
      } catch (Exception exception) {
        log.error("Error : Cannot set SSL Config properties.", exception);
      }
  
      return props;
    }
  
    public void setOtherConfig(Properties props) {
      props.put(AdminClientConfig.RETRIES_CONFIG, adminClientProperties.getRetriesConfig());
      props.put(
          AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, adminClientProperties.getRequestTimeOutMs());
      props.put(
          AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, adminClientProperties.getRetryBackOffMsConfig());
    }
  
    public Pair<String, RestTemplate> getRequestDetails(
        String suffixUrl, KafkaSupportedProtocol protocol) {
      String connectorsUrl = "";
  
      if (PLAINTEXT == protocol) {
        connectorsUrl = HTTP_PREFIX + suffixUrl;
        if (!restTemplateMap.containsKey(PLAIN_REST_TEMPLATE)) {
          restTemplateMap.put(PLAIN_REST_TEMPLATE, new RestTemplate());
        }
        return Pair.of(connectorsUrl, restTemplateMap.get(PLAIN_REST_TEMPLATE));
      } else if (SSL == protocol) {
        connectorsUrl = HTTPS_PREFIX.concat(suffixUrl);
        if (!restTemplateMap.containsKey(SSL_REST_TEMPLATE)) {
          restTemplateMap.put(
              SSL_REST_TEMPLATE, new RestTemplate(sslContextConfig.getClientHttpRequestFactory()));
        }
        return Pair.of(connectorsUrl, restTemplateMap.get(SSL_REST_TEMPLATE));
      }
      return Pair.of(connectorsUrl, new RestTemplate());
    }
  
    public HttpHeaders createHeaders(
        String clusterIdentification, KafkaClustersType kafkaClustersType) {
      String credentials = null;
      if (KafkaClustersType.SCHEMA_REGISTRY.equals(kafkaClustersType)) {
        credentials =
            env.getProperty(
                clusterIdentification.toLowerCase().concat(KAFKA_SR_CREDENTIALS_PROPERTY_SFX));
      } else if (KafkaClustersType.KAFKA_CONNECT.equals(kafkaClustersType)) {
        credentials =
            env.getProperty(
                clusterIdentification.toLowerCase().concat(KAFKA_CONNECT_CREDENTIALS_PROPERTY_SFX));
      } else if (KafkaClustersType.KAFKA.equals(kafkaClustersType)) {
        credentials =
            env.getProperty(
                clusterIdentification
                    .toLowerCase()
                    .concat(KAFKA_CONFLUENT_CLOUD_CREDENTIALS_PROPERTY_SFX));
      }
  
      HttpHeaders httpHeaders = new HttpHeaders();
      if (credentials != null) {
        byte[] encodedAuth = Base64.encodeBase64(credentials.getBytes(), false);
        String authHeader = "Basic " + new String(encodedAuth);
        httpHeaders.set("Authorization", authHeader);
      }
  
      return httpHeaders;
    }
  }
  