package se.vgregion.ssl;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.KeyStore;
import java.security.NoSuchAlgorithmException;

/**
 * @author Patrik Bergström
 */
public class ConvenientSslContextFactory {

    private String trustStore;
    private String trustStorePassword;
    private String keyStore;
    private String keyStorePassword;

    /**
     * Constructor.
     *
     * @param trustStore the trustStore either as a classpath resource name (see
     * {@link ClassLoader#getResourceAsStream(java.lang.String)} or a file system path (see
     * {@link java.io.FileInputStream#FileInputStream(java.lang.String)}), or <code>null</code>
     * @param trustStorePassword the trustStore password, or <code>null</code>
     * @param keyStore the keyStore either as a classpath resource name (see
     * {@link ClassLoader#getResourceAsStream(java.lang.String)} or a file system path (see
     * {@link java.io.FileInputStream#FileInputStream(java.lang.String)}), or <code>null</code>
     * @param keyStorePassword the keyStore password, or <code>null</code>
     */
    public ConvenientSslContextFactory(String trustStore, String trustStorePassword, String keyStore,
                                       String keyStorePassword) {
        this.trustStore = trustStore;
        this.trustStorePassword = trustStorePassword;
        this.keyStore = keyStore;
        this.keyStorePassword = keyStorePassword;
    }

    public SSLContext createSslContext() throws Exception {
        try {
            SSLContext sslContext = SSLContext.getInstance("TLS");
            sslContext.init(getKeyManagers(), getTrustManagers(), null);
            return sslContext;
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Get {@link javax.net.ssl.TrustManager} array.
     *
     * @return Array of {@link javax.net.ssl.TrustManager}s.
     * @throws Exception Exception
     */
    public TrustManager[] getTrustManagers() {
        TrustManager[] trustStoreManagers;

        InputStream tsStream = null;
        try {
            KeyStore trustedCertStore = KeyStore.getInstance("jks");
            tsStream = getClass().getClassLoader().getResourceAsStream(trustStore);

            if (tsStream == null) {
                //try absolute location on disk
                tsStream = new FileInputStream(trustStore);
                if (tsStream == null) {
                    throw new RuntimeException("Could not find truststore " + trustStore);
                }
            }

            trustedCertStore.load(tsStream, trustStorePassword.toCharArray());

            TrustManagerFactory tmf =
                    TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());

            tmf.init(trustedCertStore);
            trustStoreManagers = tmf.getTrustManagers();
            return trustStoreManagers;
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            if (tsStream != null) {
                try {
                    tsStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * Get {@link javax.net.ssl.KeyManager} array.
     *
     * @return Array of {@link javax.net.ssl.KeyManager}s.
     * @throws Exception Exception
     */
    public KeyManager[] getKeyManagers() {
        ByteArrayInputStream bin = null;
        try {
            KeyManagerFactory kmf =
                    KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
            KeyStore ks = KeyStore.getInstance("jks");
            KeyManager[] keystoreManagers = null;

            byte[] sslCert = loadClientCredential(keyStore);

            if (sslCert != null && sslCert.length > 0) {
                bin = new ByteArrayInputStream(sslCert);
                ks.load(bin, keyStorePassword.toCharArray());
                kmf.init(ks, keyStorePassword.toCharArray());
                keystoreManagers = kmf.getKeyManagers();
            }
            return keystoreManagers;
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            if (bin != null) {
                try {
                    bin.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    private byte[] loadClientCredential(String fileName) throws IOException {
        if (fileName == null) {
            return new byte[0];
        }

        InputStream in = null;
        try {
            in = getClass().getClassLoader().getResourceAsStream(fileName);

            if (in == null) {
                //try absolute location on disk
                in = new FileInputStream(fileName);
            }

            ByteArrayOutputStream out = new ByteArrayOutputStream();
            final int i1 = 512;
            byte[] buf = new byte[i1];
            int i = in.read(buf);
            while (i > 0) {
                out.write(buf, 0, i);
                i = in.read(buf);
            }
            return out.toByteArray();
        } finally {
            if (in != null) {
                in.close();
            }
        }
    }

    public String getTrustStore() {
        return trustStore;
    }

    public void setTrustStore(String trustStore) {
        this.trustStore = trustStore;
    }

    public String getTrustStorePassword() {
        return trustStorePassword;
    }

    public void setTrustStorePassword(String trustStorePassword) {
        this.trustStorePassword = trustStorePassword;
    }

    public String getKeyStore() {
        return keyStore;
    }

    public void setKeyStore(String keyStore) {
        this.keyStore = keyStore;
    }

    public String getKeyStorePassword() {
        return keyStorePassword;
    }

    public void setKeyStorePassword(String keyStorePassword) {
        this.keyStorePassword = keyStorePassword;
    }
}

@@ -683,16 +683,86 @@
 ssword;%0D
+%0A    private String trustStoreType;%0D%0A    private String keyStoreType;%0D
 %0A%0D%0A    /
@@ -1840,32 +1840,683 @@
 yStorePassword;%0D
+%0A%0D%0A        // For backwards compatibility%0D%0A        this.trustStoreType = %22jks%22;%0D%0A        this.keyStoreType = %22jks%22;%0D%0A    %7D%0D%0A%0D%0A    public ConvenientSslContextFactory(String trustStore, String trustStorePassword, String trustStoreType, String keyStore,%0D%0A                                       String keyStorePassword, String keyStoreType) %7B%0D%0A        this.trustStore = trustStore;%0D%0A        this.trustStorePassword = trustStorePassword;%0D%0A        this.keyStore = keyStore;%0D%0A        this.keyStorePassword = keyStorePassword;%0D%0A%0D%0A        // For backwards compatibility%0D%0A        this.trustStoreType = trustStoreType;%0D%0A        this.keyStoreType = keyStoreType;%0D
 %0A    %7D%0D%0A%0D%0A    pu
@@ -3249,37 +3249,46 @@
 ore.getInstance(
-%22jks%22
+trustStoreType
 );%0D%0A            
@@ -4825,21 +4825,28 @@
 nstance(
-%22jks%22
+keyStoreType
 );%0D%0A    

Add possibility to choose keystore and truststore type.