package org.jgroups.util;


import org.jgroups.logging.Log;
import org.jgroups.logging.LogFactory;

import javax.net.ssl.*;
import java.io.*;
import java.security.GeneralSecurityException;
import java.security.KeyStore;

/**
 * SslContextFactory.
 *
 * @author Tristan Tarrant
 */
public class SSLContextFactory {
    public static final String DEFAULT_SSL_PROTOCOL = "TLSv1.2";
    protected static final Log log = LogFactory.getLog(SSLContextFactory.class);
    private static final String DEFAULT_KEYSTORE_TYPE = "JKS";
    private static final String CLASSPATH_RESOURCE = "classpath:";
    private static       String SSL_PROTOCOL_PREFIX="";

    private KeyStore keyStore;
    private String keyStoreFileName;
    private char[] keyStorePassword;
    private char[] keyStoreCertificatePassword;
    private String keyStoreType = DEFAULT_KEYSTORE_TYPE;
    private String keyAlias;
    private String trustStoreFileName;
    private char[] trustStorePassword;
    private String trustStoreType = DEFAULT_KEYSTORE_TYPE;
    private String sslProtocol = DEFAULT_SSL_PROTOCOL;
    private boolean useNativeIfAvailable = true;
    private ClassLoader classLoader;

    public SSLContextFactory() {
    }

    public SSLContextFactory keyStore(KeyStore keyStore) {
        this.keyStore = keyStore;
        return this;
    }

    public SSLContextFactory keyStoreFileName(String keyStoreFileName) {
        this.keyStoreFileName = keyStoreFileName;
        return this;
    }

    public SSLContextFactory keyStorePassword(char[] keyStorePassword) {
        this.keyStorePassword = keyStorePassword;
        return this;
    }

    public SSLContextFactory keyStoreCertificatePassword(char[] keyStoreCertificatePassword) {
        this.keyStoreCertificatePassword = keyStoreCertificatePassword;
        return this;
    }

    public SSLContextFactory keyStoreType(String keyStoreType) {
        if (keyStoreType != null) {
            this.keyStoreType = keyStoreType;
        }
        return this;
    }

    public SSLContextFactory keyAlias(String keyAlias) {
        this.keyAlias = keyAlias;
        return this;
    }

    public SSLContextFactory trustStoreFileName(String trustStoreFileName) {
        this.trustStoreFileName = trustStoreFileName;
        return this;
    }

    public SSLContextFactory trustStorePassword(char[] trustStorePassword) {
        this.trustStorePassword = trustStorePassword;
        return this;
    }

    public SSLContextFactory trustStoreType(String trustStoreType) {
        if (trustStoreType != null) {
            this.trustStoreType = trustStoreType;
        }
        return this;
    }

    public SSLContextFactory sslProtocol(String sslProtocol) {
        if (sslProtocol != null) {
            this.sslProtocol = sslProtocol;
        }
        return this;
    }

    public SSLContextFactory useNativeIfAvailable(boolean useNativeIfAvailable) {
        this.useNativeIfAvailable = useNativeIfAvailable;
        return this;
    }

    public SSLContextFactory classLoader(ClassLoader classLoader) {
        this.classLoader = classLoader;
        return this;
    }

    public SSLContext getContext() {
        if(useNativeIfAvailable)
            initNative();
        try {
            KeyManager[] keyManagers = null;
            if (keyStoreFileName != null) {
                KeyManagerFactory kmf = getKeyManagerFactory();
                keyManagers = kmf.getKeyManagers();
            }
            TrustManager[] trustManagers = null;
            if (trustStoreFileName != null) {
                TrustManagerFactory tmf = getTrustManagerFactory();
                trustManagers = tmf.getTrustManagers();
            }
            String protocol = useNativeIfAvailable ? SSL_PROTOCOL_PREFIX + sslProtocol : sslProtocol;
            SSLContext sslContext = SSLContext.getInstance(protocol);
            sslContext.init(keyManagers, trustManagers, null);
            return sslContext;
        } catch (Exception e) {
            System.err.printf("ERROR %s\n", e.getMessage());
            e.printStackTrace(System.err);
            throw new RuntimeException("Cannot initialize SSL", e);
        }
    }

    public KeyManagerFactory getKeyManagerFactory() throws IOException, GeneralSecurityException {
        KeyStore ks;
        if (keyStore != null) {
            ks = keyStore;
        } else {
            ks = KeyStore.getInstance(keyStoreType != null ? keyStoreType : DEFAULT_KEYSTORE_TYPE);
            loadKeyStore(ks, keyStoreFileName, keyStorePassword, classLoader);
        }
        char[] keyPassword = keyStoreCertificatePassword == null ? keyStorePassword : keyStoreCertificatePassword;
        if (keyAlias != null) {
            if (ks.containsAlias(keyAlias) && ks.isKeyEntry(keyAlias)) {
                KeyStore.PasswordProtection passParam = new KeyStore.PasswordProtection(keyPassword);
                KeyStore.Entry entry = ks.getEntry(keyAlias, passParam);
                // Recreate the keystore with just one key
                ks = KeyStore.getInstance(keyStoreType != null ? keyStoreType : DEFAULT_KEYSTORE_TYPE);
                ks.load(null);
                ks.setEntry(keyAlias, entry, passParam);
            } else {
                throw new RuntimeException("No alias '" + keyAlias + "' in keystore '" + keyStoreFileName);
            }
        }
        KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
        kmf.init(ks, keyPassword);
        return kmf;
    }

    public TrustManagerFactory getTrustManagerFactory() throws IOException, GeneralSecurityException {
        KeyStore ks;
        if (keyStore != null) {
            ks = keyStore;
        } else {
            ks = KeyStore.getInstance(trustStoreType != null ? trustStoreType : DEFAULT_KEYSTORE_TYPE);
            loadKeyStore(ks, trustStoreFileName, trustStorePassword, classLoader);
        }
        TrustManagerFactory tmf = TrustManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
        tmf.init(ks);
        return tmf;
    }


    public static SSLEngine getEngine(SSLContext sslContext, boolean useClientMode, boolean needClientAuth) {
        SSLEngine sslEngine = sslContext.createSSLEngine();
        sslEngine.setUseClientMode(useClientMode);
        sslEngine.setNeedClientAuth(needClientAuth);
        return sslEngine;
    }

    private SSLContextFactory initNative() {
        try {
            org.wildfly.openssl.OpenSSLProvider.register();
            org.wildfly.openssl.SSL.getInstance();
            SSL_PROTOCOL_PREFIX = "openssl.";
            log.info("Using OpenSSL");
        } catch (Throwable e) {
            log.info("Using JDK SSL");
            useNativeIfAvailable=false;
        }
        return this;
    }

    private static void loadKeyStore(KeyStore ks, String keyStoreFileName, char[] keyStorePassword, ClassLoader classLoader) throws IOException, GeneralSecurityException {
        InputStream is = null;
        try {
            if (keyStoreFileName.startsWith(CLASSPATH_RESOURCE)) {
                String fileName = keyStoreFileName.substring(keyStoreFileName.indexOf(':') + 1);
                is = Util.getResourceAsStream(fileName, classLoader);
                if (is == null) {
                    throw new RuntimeException("Cannot find '" + keyStoreFileName + "' in classpath");
                }
            } else {
                File keyStoreFile = new File(keyStoreFileName);
                if (keyStoreFile.exists()) {
                    is = new FileInputStream(keyStoreFileName);
                } else {
                    is = Util.getResourceAsStream(keyStoreFileName, classLoader);
                }
            }
            ks.load(new BufferedInputStream(is), keyStorePassword);
        } finally {
            Util.close(is);
        }
    }
}
