/*
 * 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.drill.exec.server.rest.ssl;

import org.apache.commons.lang3.RandomStringUtils;
import org.apache.drill.common.config.DrillConfig;
import org.apache.drill.exec.ExecConstants;
import org.apache.drill.exec.ssl.SSLConfig;
import org.apache.drill.exec.ssl.SSLConfigBuilder;
import org.bouncycastle.asn1.x500.X500NameBuilder;
import org.bouncycastle.asn1.x500.style.BCStyle;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.jcajce.JcaX509v3CertificateBuilder;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.eclipse.jetty.util.ssl.SslContextFactory;
import org.joda.time.DateTime;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.math.BigInteger;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.util.Date;
import java.util.List;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * Configures {@link SslContextFactory} when https is enabled for Web UI
 */
public class SslContextFactoryConfigurator {
  private static final Logger logger = LoggerFactory.getLogger(SslContextFactoryConfigurator.class);

  private final DrillConfig config;
  private final String drillbitEndpointAddress;

  public SslContextFactoryConfigurator(DrillConfig config, String drillbitEndpointAddress) {
    this.config = config;
    this.drillbitEndpointAddress = drillbitEndpointAddress;
  }

  /**
   * Tries to apply ssl options configured by user. If provided configuration isn't valid,
   * new self-signed certificate will be generated and used in sslContextFactory.
   *
   * @return new configured sslContextFactory
   * @throws Exception when generation of self-signed certificate failed
   */
  public SslContextFactory configureNewSslContextFactory() throws Exception {
    SSLConfig sslConf = new SSLConfigBuilder()
        .config(config)
        .mode(SSLConfig.Mode.SERVER)
        .initializeSSLContext(false)
        .validateKeyStore(true)
        .build();
    final SslContextFactory sslContextFactory = new SslContextFactory();
    if (sslConf.isSslValid()) {
      useOptionsConfiguredByUser(sslContextFactory, sslConf);
    } else {
      useAutoGeneratedSelfSignedCertificate(sslContextFactory);
    }
    return sslContextFactory;
  }

  private void useOptionsConfiguredByUser(SslContextFactory sslFactory, SSLConfig sslConf) {
    logger.info("Using configured SSL settings for web server");
    sslFactory.setKeyStorePath(sslConf.getKeyStorePath());
    sslFactory.setKeyStorePassword(sslConf.getKeyStorePassword());
    sslFactory.setKeyManagerPassword(sslConf.getKeyPassword());
    if (sslConf.hasTrustStorePath()) {
      sslFactory.setTrustStorePath(sslConf.getTrustStorePath());
      if (sslConf.hasTrustStorePassword()) {
        sslFactory.setTrustStorePassword(sslConf.getTrustStorePassword());
      }
    }
    sslFactory.setProtocol(sslConf.getProtocol());
    sslFactory.setIncludeProtocols(sslConf.getProtocol());
    logger.info("Web server configured to use TLS protocol '{}'", sslConf.getProtocol());
    if (config.hasPath(ExecConstants.HTTP_JETTY_SSL_CONTEXT_FACTORY_OPTIONS_PREFIX)) {
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_CERT_ALIAS, sslFactory::setCertAlias);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_CRL_PATH, sslFactory::setCrlPath);
      setBooleanIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_ENABLE_CRLDP, sslFactory::setEnableCRLDP);
      setBooleanIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_ENABLE_OCSP, sslFactory::setEnableOCSP);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_ENDPOINT_IDENTIFICATION_ALGORITHM, sslFactory::setEndpointIdentificationAlgorithm);
      setStringArrayIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_EXCLUDE_CIPHER_SUITES, sslFactory::setExcludeCipherSuites);
      setStringArrayIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_EXCLUDE_PROTOCOLS, sslFactory::setExcludeProtocols);
      setStringArrayIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_INCLUDE_CIPHER_SUITES, sslFactory::setIncludeCipherSuites);
      setStringArrayIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_INCLUDE_PROTOCOLS, sslFactory::setIncludeProtocols);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_KEY_MANAGER_FACTORY_ALGORITHM, sslFactory::setKeyManagerFactoryAlgorithm);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_KEYSTORE_PROVIDER, sslFactory::setKeyStoreProvider);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_KEYSTORE_TYPE, sslFactory::setKeyStoreType);
      setIntIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_MAX_CERT_PATH_LENGTH, sslFactory::setMaxCertPathLength);
      setBooleanIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_NEED_CLIENT_AUTH, sslFactory::setNeedClientAuth);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_OCSP_RESPONDER_URL, sslFactory::setOcspResponderURL);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_PROVIDER, sslFactory::setProvider);
      setBooleanIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_RENEGOTIATION_ALLOWED, sslFactory::setRenegotiationAllowed);
      setIntIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_RENEGOTIATION_LIMIT, sslFactory::setRenegotiationLimit);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_SECURE_RANDOM_ALGORITHM, sslFactory::setSecureRandomAlgorithm);
      setBooleanIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_SESSION_CACHING_ENABLED, sslFactory::setSessionCachingEnabled);
      setIntIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_SSL_SESSION_CACHE_SIZE, sslFactory::setSslSessionCacheSize);
      setIntIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_SSL_SESSION_TIMEOUT, sslFactory::setSslSessionTimeout);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_TRUSTMANAGERFACTORY_ALGORITHM, sslFactory::setTrustManagerFactoryAlgorithm);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_TRUSTSTORE_PROVIDER, sslFactory::setTrustStoreProvider);
      setStringIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_TRUSTSTORE_TYPE, sslFactory::setTrustStoreType);
      setBooleanIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_USE_CIPHER_SUITE_ORDER, sslFactory::setUseCipherSuitesOrder);
      setBooleanIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_VALIDATE_CERTS, sslFactory::setValidateCerts);
      setBooleanIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_VALIDATE_PEER_CERTS, sslFactory::setValidatePeerCerts);
      setBooleanIfPresent(ExecConstants.HTTP_JETTY_SERVER_SSL_CONTEXT_FACTORY_WANT_CLIENT_AUTH, sslFactory::setWantClientAuth);
    }
  }

  private void setStringArrayIfPresent(String optKey, Consumer<String[]> optSet) {
    setIfPresent(optKey,
        key -> {
          List<String> list = config.getStringList(key);
          return list == null ? null : list.toArray(new String[0]);
        },
        optSet);
  }

  private void setBooleanIfPresent(String optKey, Consumer<Boolean> optSet) {
    setIfPresent(optKey, config::getBoolean, optSet);
  }

  private void setStringIfPresent(String optKey, Consumer<String> optSet) {
    setIfPresent(optKey, config::getString, optSet);
  }

  private void setIntIfPresent(String optKey, Consumer<Integer> optSet) {
    setIfPresent(optKey, config::getInt, optSet);
  }

  private <T> void setIfPresent(String optKey, Function<String, T> optGet, Consumer<T> optSet) {
    if (config.hasPath(optKey)) {
      T optVal = optGet.apply(optKey);
      if (optVal != null) {
        optSet.accept(optVal);
      }
    }
  }


  private void useAutoGeneratedSelfSignedCertificate(SslContextFactory sslContextFactory) throws Exception {
    logger.info("Using generated self-signed SSL settings for web server");
    final SecureRandom random = new SecureRandom();

    // Generate a private-public key pair
    final KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
    keyPairGenerator.initialize(1024, random);
    final KeyPair keyPair = keyPairGenerator.generateKeyPair();


    // Create builder for certificate attributes
    final X500NameBuilder nameBuilder = new X500NameBuilder(BCStyle.INSTANCE)
        .addRDN(BCStyle.OU, "Apache Drill (auth-generated)")
        .addRDN(BCStyle.O, "Apache Software Foundation (auto-generated)")
        .addRDN(BCStyle.CN, drillbitEndpointAddress);

    final DateTime now = DateTime.now();
    final Date notBefore = now.minusMinutes(1).toDate();
    final Date notAfter = now.plusYears(5).toDate();
    final BigInteger serialNumber = new BigInteger(128, random);

    // Create a certificate valid for 5years from now.
    final X509v3CertificateBuilder certificateBuilder = new JcaX509v3CertificateBuilder(
        nameBuilder.build(), // attributes
        serialNumber,
        notBefore,
        notAfter,
        nameBuilder.build(),
        keyPair.getPublic());

    // Sign the certificate using the private key
    final ContentSigner contentSigner =
        new JcaContentSignerBuilder("SHA256WithRSAEncryption").build(keyPair.getPrivate());
    final X509Certificate certificate =
        new JcaX509CertificateConverter().getCertificate(certificateBuilder.build(contentSigner));

    // Check the validity
    certificate.checkValidity(now.toDate());

    // Make sure the certificate is self-signed.
    certificate.verify(certificate.getPublicKey());

    // Generate a random password for keystore protection
    final String keyStorePasswd = RandomStringUtils.random(20);
    final KeyStore keyStore = KeyStore.getInstance("JKS");
    keyStore.load(null, null);
    keyStore.setKeyEntry("DrillAutoGeneratedCert", keyPair.getPrivate(),
        keyStorePasswd.toCharArray(), new java.security.cert.Certificate[]{certificate});

    sslContextFactory.setKeyStore(keyStore);
    sslContextFactory.setKeyStorePassword(keyStorePasswd);
  }
}