package com.jscape.util.net.ssl;

import com.jscape.util.ArrayTools;
import com.jscape.util.SuppressObfuscation;
import com.jscape.util.crypto.CryptoServicesProvider;
import com.jscape.util.reflection.FieldGetter;
import com.jscape.util.reflection.MethodInvoker;
import com.jscape.util.reflection.ObjectCreator;
import com.jscape.util.reflection.ReflectionException;

import javax.net.ssl.*;
import java.security.*;
import java.security.cert.Certificate;
import java.util.*;
import java.util.concurrent.Callable;

import static com.jscape.util.Parameters.assertNotNull;

/**
 * @author Alex Usun
 */
@SuppressObfuscation
public class SslContextFactory {

    public static SSLContext contextFor(Provider jsseProvider, final KeyManager[] keyManagers, final TrustManager[] trustManagers, final SecureRandom secureRandom, boolean fipsModeRequired)
            throws Exception {
        final SSLContext context = jsseProvider != null ? SSLContext.getInstance("TLSv1.2", jsseProvider) : SSLContext.getInstance("TLSv1.2");

        if (!fipsModeRequired) {
            context.init(keyManagers, trustManagers, secureRandom);
        } else {
            CryptoServicesProvider.applyFipsMode((Callable<Void>) () -> {
                context.init(keyManagers, trustManagers, secureRandom);
                return null;
            });
        }

        return context;
    }

    public static SSLContext contextFor(KeyManager[] keyManagers, TrustManager[] trustManagers, boolean fipsModeRequired)
            throws Exception {
        return contextFor(CryptoServicesProvider.jsseProvider(fipsModeRequired), keyManagers, trustManagers, CryptoServicesProvider.secureRandom(fipsModeRequired), fipsModeRequired);
    }

    public static SSLContext contextFor(PrivateKey key, Certificate[] keyCertificates, Certificate[] trustedCertificates, boolean fipsModeRequired)
            throws Exception {
        Provider jsseProvider = CryptoServicesProvider.jsseProvider(fipsModeRequired);
        
        KeyStore keyStore = CryptoServicesProvider.keyStoreFor(fipsModeRequired);
        addKey(keyStore, key, keyCertificates);
        addTrustedCertificates(keyStore, trustedCertificates);
        
        KeyManager[] keyManagers = keyManagersFor(jsseProvider, keyStore, fipsModeRequired);
        TrustManager[] trustManagers = trustManagersFor(jsseProvider, keyStore, fipsModeRequired);
        SecureRandom secureRandom = CryptoServicesProvider.secureRandom(fipsModeRequired);

        return contextFor(jsseProvider, keyManagers, trustManagers, secureRandom, fipsModeRequired);
    }

    public static SSLContext contextFor(PrivateKey key, Certificate[] keyCertificates, Certificate[] trustedCertificates)
            throws Exception {
        return contextFor(key, keyCertificates, trustedCertificates, CryptoServicesProvider.currentThreadFipsModeEnabled());
    }

    public static SSLContext contextFor(PrivateKey key, Certificate[] keyCertificates, TrustManager trustManager, boolean fipsModeRequired)
            throws Exception {
        Provider jsseProvider = CryptoServicesProvider.jsseProvider(fipsModeRequired);
    
        KeyStore keyStore = CryptoServicesProvider.keyStoreFor(fipsModeRequired);
        addKey(keyStore, key, keyCertificates);
        
        KeyManager[] keyManagers = keyManagersFor(jsseProvider, keyStore, fipsModeRequired);
        TrustManager[] trustManagers = { trustManager };
        SecureRandom secureRandom = CryptoServicesProvider.secureRandom(fipsModeRequired);

        return contextFor(jsseProvider, keyManagers, trustManagers, secureRandom, fipsModeRequired);
    }

    public static SSLContext contextFor(PrivateKey key, Certificate[] keyCertificates, TrustManager trustManager)
            throws Exception {
        return contextFor(key, keyCertificates, trustManager, CryptoServicesProvider.currentThreadFipsModeEnabled());
    }

    public static SSLContext contextFor(PrivateKey key, Certificate[] keyCertificates, boolean fipsModeRequired)
            throws Exception {
        return contextFor(key, keyCertificates, new NullTrustManager(), fipsModeRequired);
    }

    public static SSLContext contextFor(PrivateKey key, Certificate[] keyCertificates)
            throws Exception {
        return contextFor(key, keyCertificates, new NullTrustManager());
    }
    
    public static SSLContext contextFor(Certificate[] trustedCertificates, boolean fipsModeRequired)
            throws Exception {
        Provider jsseProvider = CryptoServicesProvider.jsseProvider(fipsModeRequired);
        
        KeyStore keyStore = CryptoServicesProvider.keyStoreFor(fipsModeRequired);
        addTrustedCertificates(keyStore, trustedCertificates);
        
        TrustManager[] trustManagers = trustManagersFor(jsseProvider, keyStore, fipsModeRequired);
        SecureRandom secureRandom = CryptoServicesProvider.secureRandom(fipsModeRequired);
        
        return contextFor(jsseProvider, null, trustManagers, secureRandom, fipsModeRequired);
    }
    
    public static SSLContext contextFor(Certificate[] trustedCertificates)
            throws Exception {
        return contextFor(trustedCertificates, CryptoServicesProvider.currentThreadFipsModeEnabled());
    }

    public static SSLContext contextFor(KeyStore keyStore, String keyPassword, KeyStore trustedStore, boolean fipsModeRequired)
            throws Exception {
        PrivateKey key = null;
        Certificate[] keyCertificates = null;
        for (Enumeration<String> e = keyStore.aliases(); e.hasMoreElements();) {
            String alias = e.nextElement();
            if (keyStore.isKeyEntry(alias)) {
                key = (PrivateKey) keyStore.getKey(alias, keyPassword.toCharArray());
                keyCertificates = keyStore.getCertificateChain(alias);
                break;
            }
        }

        assertNotNull(key);
        assertNotNull(keyCertificates);

        List<Certificate> trustedCertificates = new LinkedList<>();
        for (Enumeration<String> e = trustedStore.aliases(); e.hasMoreElements();) {
            String alias = e.nextElement();
            if (trustedStore.isCertificateEntry(alias)) {
                Certificate certificate = trustedStore.getCertificate(alias);
                trustedCertificates.add(certificate);
            }
        }

        return contextFor(key, keyCertificates, trustedCertificates.toArray(new Certificate[0]), fipsModeRequired);
    }

    public static SSLContext contextFor(KeyStore keyStore, String keyPassword, KeyStore trustedStore)
            throws Exception {
        return contextFor(keyStore, keyPassword, trustedStore, CryptoServicesProvider.currentThreadFipsModeEnabled());
    }

    public static SSLContext contextFor(KeyStore keyStore, String keyPassword, boolean fipsModeRequired)
            throws Exception {
        PrivateKey key = null;
        Certificate[] keyCertificates = null;
        for (Enumeration<String> e = keyStore.aliases(); e.hasMoreElements();) {
            String alias = e.nextElement();
            if (keyStore.isKeyEntry(alias)) {
                key = (PrivateKey) keyStore.getKey(alias, keyPassword.toCharArray());
                keyCertificates = keyStore.getCertificateChain(alias);
                break;
            }
        }

        assertNotNull(key);
        assertNotNull(keyCertificates);

        return contextFor(key, keyCertificates, fipsModeRequired);
    }

    public static SSLContext contextFor(KeyStore keyStore, String keyPassword)
            throws Exception {
        return contextFor(keyStore, keyPassword, CryptoServicesProvider.currentThreadFipsModeEnabled());
    }

    public static SSLContext defaultContext(boolean fipsModeRequired)
            throws Exception {
        Provider jsseProvider = CryptoServicesProvider.jsseProvider(fipsModeRequired);
        TrustManager[] trustManagers = { new NullTrustManager() };
        SecureRandom secureRandom = CryptoServicesProvider.secureRandom(fipsModeRequired);

        return contextFor(jsseProvider, null, trustManagers, secureRandom, fipsModeRequired);
    }

    public static SSLContext defaultContext()
            throws Exception {
        return defaultContext(CryptoServicesProvider.currentThreadFipsModeEnabled());
    }

    public static SSLContext systemContextFor(KeyManager[] keyManagers, TrustManager[] trustManagers)
            throws Exception {
        SSLContext context = SSLContext.getInstance("TLSv1.2");
        context.init(keyManagers, trustManagers, new SecureRandom());

        return context;
    }

    public static SSLContext systemContextFor(KeyStore keyStore, String keyPassword, KeyStore trustedStore)
            throws Exception {
        PrivateKey key = null;
        Certificate[] keyCertificates = null;
        for (Enumeration<String> e = keyStore.aliases(); e.hasMoreElements();) {
            String alias = e.nextElement();
            if (keyStore.isKeyEntry(alias)) {
                key = (PrivateKey) keyStore.getKey(alias, keyPassword.toCharArray());
                keyCertificates = keyStore.getCertificateChain(alias);
                break;
            }
        }

        assertNotNull(key);
        assertNotNull(keyCertificates);

        List<Certificate> trustedCertificates = new LinkedList<>();
        for (Enumeration<String> e = trustedStore.aliases(); e.hasMoreElements();) {
            String alias = e.nextElement();
            if (trustedStore.isCertificateEntry(alias)) {
                Certificate certificate = trustedStore.getCertificate(alias);
                trustedCertificates.add(certificate);
            }
        }

        return systemContextFor(key, keyCertificates, trustedCertificates.toArray(new Certificate[0]));
    }

    public static SSLContext systemContextFor(KeyStore keyStore, String keyPassword)
            throws Exception {
        PrivateKey key = null;
        Certificate[] keyCertificates = null;
        for (Enumeration<String> e = keyStore.aliases(); e.hasMoreElements();) {
            String alias = e.nextElement();
            if (keyStore.isKeyEntry(alias)) {
                key = (PrivateKey) keyStore.getKey(alias, keyPassword.toCharArray());
                keyCertificates = keyStore.getCertificateChain(alias);
                break;
            }
        }

        assertNotNull(key);
        assertNotNull(keyCertificates);

        return systemContextFor(key, keyCertificates);
    }

    public static SSLContext systemContextFor(PrivateKey key, Certificate[] keyCertificates, Certificate[] trustedCertificates)
            throws Exception {
        KeyStore keyStore = KeyStore.getInstance("JKS");
        keyStore.load(null, null);
    
        addKey(keyStore, key, keyCertificates);
        addTrustedCertificates(keyStore, trustedCertificates);

        KeyManager[] keyManagers = keyManagersFor(null, keyStore, false);
        TrustManager[] trustManagers = trustManagersFor(null, keyStore, false);

        return systemContextFor(keyManagers, trustManagers);
    }
    
    public static SSLContext systemContextFor(PrivateKey key, Certificate[] keyCertificates, TrustManager trustManager)
            throws Exception {
        KeyStore keyStore = KeyStore.getInstance("JKS");
        keyStore.load(null, null);
    
        addKey(keyStore, key, keyCertificates);
    
        KeyManager[] keyManagers = keyManagersFor(null, keyStore, false);
        TrustManager[] trustManagers = { trustManager };

        return systemContextFor(keyManagers, trustManagers);
    }

    public static SSLContext systemContextFor(PrivateKey key, Certificate[] keyCertificates)
            throws Exception {
        return systemContextFor(key, keyCertificates, new NullTrustManager());
    }

    public static SSLContext systemContext()
            throws Exception {
        TrustManager[] trustManagers = { new NullTrustManager() };
        SecureRandom secureRandom = new SecureRandom();

        return contextFor(null, null, trustManagers, secureRandom, false);
    }

    public static String[] supportedProtocols(boolean fipsModeRequired) {
        try {
            SSLContext context = defaultContext(fipsModeRequired);
            return context.getSupportedSSLParameters().getProtocols();
        } catch (Exception e) {
            return new String[0];
        }
    }

    public static String[] supportedProtocols() {
        return supportedProtocols(CryptoServicesProvider.currentThreadFipsModeEnabled());
    }

    public static String[] supportedCipherSuites(boolean fipsModeRequired) {
        try {
            SSLContext context = defaultContext(fipsModeRequired);
            String[] suites = context.getSupportedSSLParameters().getCipherSuites();

            return ArrayTools.sort(suites, new PerfectForwardSecrecyCipherSuiteComparator());
        } catch (Exception e) {
            return new String[0];
        }
    }

    public static String[] supportedCipherSuites() {
        return supportedCipherSuites(CryptoServicesProvider.currentThreadFipsModeEnabled());
    }

    public static SSLSocket reuseSession(SSLSocket source, SSLSocket destination)
            throws Exception {
        SSLSession session = source.getSession();

        try {
            return reuseSessionOracle(session, destination);
        } catch (Exception e) {
            return reuseSessionBc(session, destination);
        }
    }

    private static SSLSocket reuseSessionOracle(SSLSession session, SSLSocket destination)
            throws Exception {
        SSLSessionContext sessionContext = session.getSessionContext();

        try {
            Object sessionCache = FieldGetter.get(sessionContext, "sessionHostPortCache");
            MethodInvoker.invoke(sessionCache, "put", String.format("%s:%s", destination.getInetAddress().getHostName(), destination.getPort()).toLowerCase(Locale.ROOT), session);
            MethodInvoker.invoke(sessionCache, "put", String.format("%s:%s", destination.getInetAddress().getHostAddress(), destination.getPort()).toLowerCase(Locale.ROOT), session);
        } catch (ReflectionException e) {
            Object sessionCache = FieldGetter.get(sessionContext, "sessionsByPeer");
            Object entry = MethodInvoker.invoke(sessionCache, "get", String.format("%s:%s", session.getPeerHost(), session.getPeerPort()).toLowerCase(Locale.ROOT));
            MethodInvoker.invoke(sessionCache, "put", String.format("%s:%s", destination.getInetAddress().getHostName(), destination.getPort()).toLowerCase(Locale.ROOT), entry);
            MethodInvoker.invoke(sessionCache, "put", String.format("%s:%s", destination.getInetAddress().getHostAddress(), destination.getPort()).toLowerCase(Locale.ROOT), entry);
        }

        return destination;
    }

    private static SSLSocket reuseSessionBc(SSLSession session, SSLSocket destination)
            throws Exception {
        SSLSessionContext sessionContext = session.getSessionContext();

        Object idSessionCache = FieldGetter.get(sessionContext, "sessionsByID");
        Object entry = MethodInvoker.invoke(idSessionCache, "get", ObjectCreator.create(Class.forName("org.bouncycastle.tls.SessionID"), new Class[]{byte[].class}, session.getId()));

        Object peerSessionCache = FieldGetter.get(sessionContext, "sessionsByPeer");
        MethodInvoker.invoke(peerSessionCache, "put", String.format("%s:%s", destination.getInetAddress().getHostName(), destination.getPort()).toLowerCase(Locale.ROOT), entry);
        MethodInvoker.invoke(peerSessionCache, "put", String.format("%s:%s", destination.getInetAddress().getHostAddress(), destination.getPort()).toLowerCase(Locale.ROOT), entry);

        return destination;
    }

    private static KeyManager[] keyManagersFor(Provider jsseProvider, final KeyStore keyStore, boolean fipsModeRequired)
            throws Exception {
        final KeyManagerFactory managerFactory = jsseProvider != null ? KeyManagerFactory.getInstance("X509", jsseProvider) : KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());

        if (!fipsModeRequired) {
            managerFactory.init(keyStore, new char[0]);
        } else {
            CryptoServicesProvider.applyFipsMode((Callable<Void>) () -> {
                managerFactory.init(keyStore, new char[0]);
                return null;
            });
        }

        return managerFactory.getKeyManagers();
    }

    private static TrustManager[] trustManagersFor(Provider jsseProvider, final KeyStore keyStore, boolean fipsModeRequired)
            throws Exception {
        final TrustManagerFactory managerFactory = jsseProvider != null ? TrustManagerFactory.getInstance("X509", jsseProvider) : TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());

        if (!fipsModeRequired) {
            managerFactory.init(keyStore);
        } else {
            CryptoServicesProvider.applyFipsMode((Callable<Void>) () -> {
                managerFactory.init(keyStore);
                return null;
            });
        }

        return managerFactory.getTrustManagers();
    }
    
    private static void addKey(KeyStore keyStore, PrivateKey key, Certificate... certificates)
            throws KeyStoreException {
        keyStore.setKeyEntry("key", key, new char[0], certificates);
    }
    
    private static void addTrustedCertificates(KeyStore keyStore, Certificate... certificates)
            throws KeyStoreException {
        for (Certificate certificate : certificates) {
            String alias = UUID.randomUUID().toString().replace("-", "");
            keyStore.setCertificateEntry(alias, certificate);
        }
    }

    private SslContextFactory() {
        // empty
    }
}
