package at.gv.egiz.eaaf.core.impl.http;

import java.net.Socket;
import java.security.KeyManagementException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.Principal;
import java.security.PrivateKey;
import java.security.Provider;
import java.security.SecureRandom;
import java.security.Security;
import java.security.UnrecoverableKeyException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509ExtendedKeyManager;
import javax.net.ssl.X509TrustManager;

import org.apache.http.ssl.PrivateKeyDetails;
import org.apache.http.ssl.PrivateKeyStrategy;
import org.apache.http.ssl.SSLContextBuilder;
import org.apache.http.ssl.TrustStrategy;
import org.bouncycastle.jsse.provider.BouncyCastleJsseProvider;

/**
 * Fork of {@link SSLContextBuilder} that uses JSSE provider to get TrustManager.
 * 
 * <p>This implementation fix an incompatibility between {@link BouncyCastleJsseProvider} and JAVA JDK >= v9</p>
 *
 * @author tlenz
 *
 */
public class EaafSslContextBuilder {

  static final String TLS = "TLS";

  private String protocol;
  private final Set<KeyManager> keyManagers;
  private String keyManagerFactoryAlgorithm = KeyManagerFactory.getDefaultAlgorithm();
  private String keyStoreType = KeyStore.getDefaultType();
  private final Set<TrustManager> trustManagers;
  private String trustManagerFactoryAlgorithm = TrustManagerFactory.getDefaultAlgorithm();
  private SecureRandom secureRandom;
  private Provider provider;

  public static EaafSslContextBuilder create() {
    return new EaafSslContextBuilder();
  }

  /**
   * Get a new SSLContext builder object.
   */
  public EaafSslContextBuilder() {
    super();
    this.keyManagers = new LinkedHashSet<>();
    this.trustManagers = new LinkedHashSet<>();
  }

  /**
   * Sets the SSLContext protocol algorithm name.
   *
   * @param protocol the SSLContext protocol algorithm name of the requested
   *                 protocol. See the SSLContext section in the <a href=
   *                 "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#SSLContext">Java
   *                 Cryptography Architecture Standard Algorithm Name
   *                 Documentation</a> for more information.
   * @return this builder
   * @see <a href=
   *      "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#SSLContext">Java
   *      Cryptography Architecture Standard Algorithm Name Documentation</a>
   * @deprecated Use {@link #setProtocol(String)}.
   */
  @Deprecated
  public EaafSslContextBuilder useProtocol(final String protocol) {
    this.protocol = protocol;
    return this;
  }

  /**
   * Sets the SSLContext protocol algorithm name.
   *
   * @param protocol the SSLContext protocol algorithm name of the requested
   *                 protocol. See the SSLContext section in the <a href=
   *                 "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#SSLContext">Java
   *                 Cryptography Architecture Standard Algorithm Name
   *                 Documentation</a> for more information.
   * @return this builder
   * @see <a href=
   *      "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#SSLContext">Java
   *      Cryptography Architecture Standard Algorithm Name Documentation</a>
   * @since 4.4.7
   */
  public EaafSslContextBuilder setProtocol(final String protocol) {
    this.protocol = protocol;
    return this;
  }

  public EaafSslContextBuilder setSecureRandom(final SecureRandom secureRandom) {
    this.secureRandom = secureRandom;
    return this;
  }

  public EaafSslContextBuilder setProvider(final Provider provider) {
    this.provider = provider;
    return this;
  }

  public EaafSslContextBuilder setProvider(final String name) {
    this.provider = Security.getProvider(name);
    return this;
  }

  /**
   * Sets the key store type.
   *
   * @param keyStoreType the SSLkey store type. See the KeyStore section in the
   *                     <a href=
   *                     "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#KeyStore">Java
   *                     Cryptography Architecture Standard Algorithm Name
   *                     Documentation</a> for more information.
   * @return this builder
   * @see <a href=
   *      "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#KeyStore">Java
   *      Cryptography Architecture Standard Algorithm Name Documentation</a>
   * @since 4.4.7
   */
  public EaafSslContextBuilder setKeyStoreType(final String keyStoreType) {
    this.keyStoreType = keyStoreType;
    return this;
  }

  /**
   * Sets the key manager factory algorithm name.
   *
   * @param keyManagerFactoryAlgorithm the key manager factory algorithm name of
   *                                   the requested protocol. See the
   *                                   KeyManagerFactory section in the <a href=
   *                                   "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#KeyManagerFactory">Java
   *                                   Cryptography Architecture Standard
   *                                   Algorithm Name Documentation</a> for more
   *                                   information.
   * @return this builder
   * @see <a href=
   *      "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#KeyManagerFactory">Java
   *      Cryptography Architecture Standard Algorithm Name Documentation</a>
   * @since 4.4.7
   */
  public EaafSslContextBuilder setKeyManagerFactoryAlgorithm(final String keyManagerFactoryAlgorithm) {
    this.keyManagerFactoryAlgorithm = keyManagerFactoryAlgorithm;
    return this;
  }

  /**
   * Sets the trust manager factory algorithm name.
   *
   * @param trustManagerFactoryAlgorithm the trust manager algorithm name of the
   *                                     requested protocol. See the
   *                                     TrustManagerFactory section in the
   *                                     <a href=
   *                                     "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#TrustManagerFactory">Java
   *                                     Cryptography Architecture Standard
   *                                     Algorithm Name Documentation</a> for more
   *                                     information.
   * @return this builder
   * @see <a href=
   *      "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#TrustManagerFactory">Java
   *      Cryptography Architecture Standard Algorithm Name Documentation</a>
   * @since 4.4.7
   */
  public EaafSslContextBuilder setTrustManagerFactoryAlgorithm(final String trustManagerFactoryAlgorithm) {
    this.trustManagerFactoryAlgorithm = trustManagerFactoryAlgorithm;
    return this;
  }

  /**
   * Load custom truststore.
   * 
   * @param truststore {@link KeyStore} if trusted certificates 
   * @param trustStrategy Trust validation strategy
   * @return {@link EaafSslContextBuilder}
   * @throws NoSuchAlgorithmException In case of an invalid TrustManager algorithm
   * @throws KeyStoreException In case of an invalid KeyStore
   */
  public EaafSslContextBuilder loadTrustMaterial(
      final KeyStore truststore,
      final TrustStrategy trustStrategy) throws NoSuchAlgorithmException, KeyStoreException {

    final String alg = trustManagerFactoryAlgorithm == null
        ? TrustManagerFactory.getDefaultAlgorithm()
        : trustManagerFactoryAlgorithm;

    final TrustManagerFactory tmfactory = provider != null
        ? TrustManagerFactory.getInstance(alg, provider)
        : TrustManagerFactory.getInstance(alg);
    tmfactory.init(truststore);
    final TrustManager[] tms = tmfactory.getTrustManagers();
    if (tms != null) {
      if (trustStrategy != null) {
        for (int i = 0; i < tms.length; i++) {
          final TrustManager tm = tms[i];
          if (tm instanceof X509TrustManager) {
            tms[i] = new TrustManagerDelegate((X509TrustManager) tm, trustStrategy);
          }
        }
      }
      Collections.addAll(this.trustManagers, tms);
    }
    return this;
  }

  public EaafSslContextBuilder loadTrustMaterial(
      final TrustStrategy trustStrategy) throws NoSuchAlgorithmException, KeyStoreException {
    return loadTrustMaterial(null, trustStrategy);
  }

  
  /**
   * Load SSL client-authentication key-material into SSL context.
   * 
   * @param keystore {@link KeyStore} for SSL client-authentication
   * @param keyPassword Password for this keystore
   * @param aliasStrategy Stategy to select keys by alias
   * @return {@link EaafSslContextBuilder}
   * @throws NoSuchAlgorithmException In case of an invalid KeyManagerFactory algorithm
   * @throws KeyStoreException In case of an invalid KeyStore
   * @throws UnrecoverableKeyException In case of a invalid Key in this KeyStore
   */
  public EaafSslContextBuilder loadKeyMaterial(
      final KeyStore keystore,
      final char[] keyPassword,
      final PrivateKeyStrategy aliasStrategy)
      throws NoSuchAlgorithmException, KeyStoreException, UnrecoverableKeyException {
    final KeyManagerFactory kmfactory = KeyManagerFactory
        .getInstance(keyManagerFactoryAlgorithm == null ? KeyManagerFactory.getDefaultAlgorithm()
            : keyManagerFactoryAlgorithm);
    kmfactory.init(keystore, keyPassword);
    final KeyManager[] kms = kmfactory.getKeyManagers();
    if (kms != null) {
      if (aliasStrategy != null) {
        for (int i = 0; i < kms.length; i++) {
          final KeyManager km = kms[i];
          if (km instanceof X509ExtendedKeyManager) {
            kms[i] = new KeyManagerDelegate((X509ExtendedKeyManager) km, aliasStrategy);
          }
        }
      }
      Collections.addAll(keyManagers, kms);
    }
    return this;
  }

  public EaafSslContextBuilder loadKeyMaterial(
      final KeyStore keystore,
      final char[] keyPassword) throws NoSuchAlgorithmException, KeyStoreException,
      UnrecoverableKeyException {
    return loadKeyMaterial(keystore, keyPassword, null);
  }

  protected void initSslContext(
      final SSLContext sslContext,
      final Collection<KeyManager> keyManagers,
      final Collection<TrustManager> trustManagers,
      final SecureRandom secureRandom) throws KeyManagementException {
    sslContext.init(
        !keyManagers.isEmpty() ? keyManagers.toArray(new KeyManager[keyManagers.size()]) : null,
        !trustManagers.isEmpty() ? trustManagers.toArray(new TrustManager[trustManagers.size()]) : null,
        secureRandom);
  }

  /**
   * Build a {@link SSLContext} from this builder.
   * 
   * @return new {@link SSLContext}
   * @throws NoSuchAlgorithmException In case of an unknown SSL protocol
   * @throws KeyManagementException In case of a key-access error
   */
  public SSLContext build() throws NoSuchAlgorithmException, KeyManagementException {
    final SSLContext sslContext;
    final String protocolStr = this.protocol != null ? this.protocol : TLS;
    if (this.provider != null) {
      sslContext = SSLContext.getInstance(protocolStr, this.provider);
    } else {
      sslContext = SSLContext.getInstance(protocolStr);
    }
    initSslContext(sslContext, keyManagers, trustManagers, secureRandom);
    return sslContext;
  }

  static class TrustManagerDelegate implements X509TrustManager {

    private final X509TrustManager trustManager;
    private final TrustStrategy trustStrategy;

    TrustManagerDelegate(final X509TrustManager trustManager, final TrustStrategy trustStrategy) {
      super();
      this.trustManager = trustManager;
      this.trustStrategy = trustStrategy;
    }

    @Override
    public void checkClientTrusted(
        final X509Certificate[] chain, final String authType) throws CertificateException {
      this.trustManager.checkClientTrusted(chain, authType);
    }

    @Override
    public void checkServerTrusted(
        final X509Certificate[] chain, final String authType) throws CertificateException {
      if (!this.trustStrategy.isTrusted(chain, authType)) {
        this.trustManager.checkServerTrusted(chain, authType);
      }
    }

    @Override
    public X509Certificate[] getAcceptedIssuers() {
      return this.trustManager.getAcceptedIssuers();
    }

  }

  static class KeyManagerDelegate extends X509ExtendedKeyManager {

    private final X509ExtendedKeyManager keyManager;
    private final PrivateKeyStrategy aliasStrategy;

    KeyManagerDelegate(final X509ExtendedKeyManager keyManager, final PrivateKeyStrategy aliasStrategy) {
      super();
      this.keyManager = keyManager;
      this.aliasStrategy = aliasStrategy;
    }

    @Override
    public String[] getClientAliases(
        final String keyType, final Principal[] issuers) {
      return this.keyManager.getClientAliases(keyType, issuers);
    }

    public Map<String, PrivateKeyDetails> getClientAliasMap(
        final String[] keyTypes, final Principal[] issuers) {
      final Map<String, PrivateKeyDetails> validAliases = new HashMap<>();
      for (final String keyType : keyTypes) {
        final String[] aliases = this.keyManager.getClientAliases(keyType, issuers);
        if (aliases != null) {
          for (final String alias : aliases) {
            validAliases.put(alias,
                new PrivateKeyDetails(keyType, this.keyManager.getCertificateChain(alias)));
          }
        }
      }
      return validAliases;
    }

    public Map<String, PrivateKeyDetails> getServerAliasMap(
        final String keyType, final Principal[] issuers) {
      final Map<String, PrivateKeyDetails> validAliases = new HashMap<>();
      final String[] aliases = this.keyManager.getServerAliases(keyType, issuers);
      if (aliases != null) {
        for (final String alias : aliases) {
          validAliases.put(alias,
              new PrivateKeyDetails(keyType, this.keyManager.getCertificateChain(alias)));
        }
      }
      return validAliases;
    }

    @Override
    public String chooseClientAlias(
        final String[] keyTypes, final Principal[] issuers, final Socket socket) {
      final Map<String, PrivateKeyDetails> validAliases = getClientAliasMap(keyTypes, issuers);
      return this.aliasStrategy.chooseAlias(validAliases, socket);
    }

    @Override
    public String[] getServerAliases(
        final String keyType, final Principal[] issuers) {
      return this.keyManager.getServerAliases(keyType, issuers);
    }

    @Override
    public String chooseServerAlias(
        final String keyType, final Principal[] issuers, final Socket socket) {
      final Map<String, PrivateKeyDetails> validAliases = getServerAliasMap(keyType, issuers);
      return this.aliasStrategy.chooseAlias(validAliases, socket);
    }

    @Override
    public X509Certificate[] getCertificateChain(final String alias) {
      return this.keyManager.getCertificateChain(alias);
    }

    @Override
    public PrivateKey getPrivateKey(final String alias) {
      return this.keyManager.getPrivateKey(alias);
    }

    @Override
    public String chooseEngineClientAlias(
        final String[] keyTypes, final Principal[] issuers, final SSLEngine sslEngine) {
      final Map<String, PrivateKeyDetails> validAliases = getClientAliasMap(keyTypes, issuers);
      return this.aliasStrategy.chooseAlias(validAliases, null);
    }

    @Override
    public String chooseEngineServerAlias(
        final String keyType, final Principal[] issuers, final SSLEngine sslEngine) {
      final Map<String, PrivateKeyDetails> validAliases = getServerAliasMap(keyType, issuers);
      return this.aliasStrategy.chooseAlias(validAliases, null);
    }

  }

  @Override
  public String toString() {
    return "[provider=" + provider + ", protocol=" + protocol + ", keyStoreType=" + keyStoreType
        + ", keyManagerFactoryAlgorithm=" + keyManagerFactoryAlgorithm + ", keyManagers=" + keyManagers
        + ", trustManagerFactoryAlgorithm=" + trustManagerFactoryAlgorithm + ", trustManagers="
        + trustManagers
        + ", secureRandom=" + secureRandom + "]";
  }
}