package be.iminds.ilabt.jfed.util;

import be.iminds.ilabt.jfed.lowlevel.TestbedInfoSource;

import be.iminds.ilabt.jfed.fedmon.webapi.service.json.*;
import org.slf4j.LoggerFactory;

import javax.inject.Singleton;
import java.io.*;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.util.*;

/**
 * JFedTrustStore
 */
@Singleton
public class JFedTrustStore {
    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(JFedTrustStore.class);

    private final Set<String> allowedServerCertificateHostnameAliases;
    private final KeyStore trustStore;

    public final List<String> addedPems = new ArrayList<>();
    private final List<Certificate> extraTrustedCertificates = new ArrayList<>();

    /** default constructor: initialise with copy of system trust store. */
    public JFedTrustStore() {
        this.trustStore = getSystemTrustStore();
        this.allowedServerCertificateHostnameAliases = new HashSet<>();
    }
    public JFedTrustStore(KeyStore trustStore) {
        try {
            this.trustStore = copyTrustStore(trustStore, null, null, "somepass".toCharArray());
        } catch (Exception e) {
            throw new RuntimeException("Could not copy trust store: "+e.getMessage(), e);
        }
        this.allowedServerCertificateHostnameAliases = new HashSet<>();
    }
    public JFedTrustStore(JFedTrustStore jFedTrustStore) {
        try {
            this.trustStore = copyTrustStore(jFedTrustStore.trustStore, null, null, "somepass".toCharArray());
        } catch (Exception e) {
            throw new RuntimeException("Could not copy trust store: "+e.getMessage(), e);
        }
        this.allowedServerCertificateHostnameAliases =
                new HashSet<>(jFedTrustStore.getAllowedServerCertificateHostnameAliases());
    }


    /**
     * Create a trust store that has certificates for the specified authority
     * Also adds the AllowedCertificateHostnameAliases of the authority
     * */
    public JFedTrustStore(Server server) {
        this.trustStore = getSystemTrustStore();
        if (server.getCertificateChain() != null)
            addTrustedPemCertificateIfNotAdded(server.getCertificateChain());
        this.allowedServerCertificateHostnameAliases = new HashSet<>();
        if (server.getAllowedCertificateAlias() != null && !server.getAllowedCertificateAlias().trim().isEmpty())
            this.allowedServerCertificateHostnameAliases.add(server.getAllowedCertificateAlias());
    }
    /**
     * Create a trust store that has certificates of all authorities in the list added
     *
     * (doesn't add the AllowedCertificateHostnameAliases of any authority)
     * */
    public JFedTrustStore(TestbedInfoSource testbedInfoSource) {
        this.trustStore = getSystemTrustStore();
        addAuthorityCerts(testbedInfoSource);
        allowedServerCertificateHostnameAliases = new HashSet<>();
    }


    public void addAuthorityCerts(TestbedInfoSource testbedInfoSource) {
        testbedInfoSource.getServers().stream()
                .filter(server -> server.getCertificateChain() != null)
                .forEach(server -> addTrustedPemCertificateIfNotAdded(server.getCertificateChain()));
    }
    public void addAuthorityCert(Server server) {
        if (server.getCertificateChain() != null)
            addTrustedPemCertificateIfNotAdded(server.getCertificateChain());
    }

    public Collection<String> getAllowedServerCertificateHostnameAliases() {
        return Collections.unmodifiableCollection(allowedServerCertificateHostnameAliases);
    }

    public KeyStore getTrustStore() {
        return trustStore;
    }

    public void addAllowedServerCertificateHostnameAlias(String alias) {
        if (alias != null)
            allowedServerCertificateHostnameAliases.add(alias);
    }

    public void addAllowedServerCertificateHostnameAliases(Collection<String> aliases) {
        if (aliases != null)
            allowedServerCertificateHostnameAliases.addAll(aliases);
    }

    public void addTrustedCertificate(Certificate c) {
        assert c != null;
        //TODO: do not add if already added
        extraTrustedCertificates.add(c);
        try {
            trustStore.setCertificateEntry("extraCert" +extraTrustedCertificates.size() , c);
        } catch (KeyStoreException e) {
            LOG.error("Error while creating adding certificate to trust store", e);
            throw new RuntimeException("Could not add certificate to trust store: " + e.getMessage(), e);
        }
    }

    public void addTrustedPemCertificate(String pem) {
        assert pem != null;
        addTrustedCertificate(KeyUtil.pemToX509Certificate(pem));
    }

    public void addTrustedPemCertificateIfNotAdded(String pem) {
        assert (pem != null);
        if (addedPems.contains(pem)) return;
        addedPems.add(pem);
        addTrustedPemCertificate(pem);
    }




    //system trust store
    private static KeyStore systemTrustStore = null;
    /**
     * Note: there is no good way to find the loaded system Trust Store, so we have to search for the Trust Store file.
     * There are default locations for that, and these can be overwritten by system properties.
     * <p/>
     * TODO: If we cannot find a file, we fall back to an internal trust store.
     *
     * @return the a KeyStore containing all certificates in the java system Trust Store.
     */
    public static synchronized KeyStore getSystemTrustStore() {
        //default truststore password can be overriden by setting property javax.net.ssl.trustStorePassword
        String systemTrustStorePass = System.getProperty("javax.net.ssl.trustStorePassword");
        if (systemTrustStorePass == null)
            systemTrustStorePass = "changeit"; //default password for system truststore

        if (systemTrustStore != null) {
            try {
                return copyTrustStore(systemTrustStore, systemTrustStorePass.toCharArray(), null, systemTrustStorePass.toCharArray());
            } catch (Exception e) {
                LOG.error("Failed to copy system trust store", e);
                throw new RuntimeException("Failed to copy system trust store", e);
            }
        }

        final String jssecacertsFilename =
                System.getProperty("java.home") + File.separator + "lib" + File.separator + "security" + File.separator + "jssecacerts";
        final String cacertsFilename =
                System.getProperty("java.home") + File.separator + "lib" + File.separator + "security" + File.separator + "cacerts";

        //default can be overriden by setting property javax.net.ssl.trustStore
        String systemTrustStoreFilename = System.getProperty("javax.net.ssl.trustStore");
        if (systemTrustStoreFilename == null) {
            //if property not set, system default is used. If jssecacerts does no exist, cacerts should be used.
            systemTrustStoreFilename = jssecacertsFilename;
            if (!new File(systemTrustStoreFilename).exists())
                systemTrustStoreFilename = cacertsFilename;
            if (!new File(systemTrustStoreFilename).exists()) {
                //TODO fall back to an internal trust store, and warn about that
                throw new RuntimeException("Could not find any system trust store!");
            }
        }

        try {
            systemTrustStore = KeyStore.getInstance(KeyStore.getDefaultType());
            FileInputStream system_trustStore_in = new FileInputStream(systemTrustStoreFilename);
            systemTrustStore.load(system_trustStore_in, systemTrustStorePass.toCharArray());
            system_trustStore_in.close();
        } catch (Exception e) {
            LOG.error("ERROR loading system trust store: " + e.getMessage(), e);
            LOG.error("  Normally, the trust store is at one of these locations:\n" +
                    "   - <JAVA_HOME>/lib/security/jssecacerts => \"" + jssecacertsFilename + "\"\n" +
                    "   - <JAVA_HOME>/lib/security/cacerts => \"" + cacertsFilename + "\"\n");
            LOG.error("  You can use another by setting the system property \"javax.net.ssl.trustStore\"");
            LOG.error("  You can specify a non default password with \"javax.net.ssl.trustStorePassword\" (default pass is \"changeit\")");
            systemTrustStore = null;
            throw new RuntimeException("Cannot locate and load system trust store: " + e.getMessage(), e);
        }

        try {
            //always return a copy to be safe. We don't want the systemTrustStore itself to be ever modified.
            return copyTrustStore(systemTrustStore, systemTrustStorePass.toCharArray(), null, systemTrustStorePass.toCharArray());
        } catch (Exception e) {
            LOG.error("Could not copy the found system trust store.", e);
            throw new RuntimeException("Could not copy the found system trust store.", e);
        }
    }



    //general helpers
    public static KeyStore copyTrustStore(KeyStore origKeyStore, char[] origPass, char[] origKeyPass, char[] resultPass)
            throws KeyStoreException, UnrecoverableKeyException, NoSuchAlgorithmException, IOException, CertificateException {
        return copyKeyStore(origKeyStore, origPass, origKeyPass, resultPass);
    }
    public static KeyStore copyKeyStore(KeyStore origKeyStore, char[] origPass, char[] origKeyPass, char[] resultPass)
            throws KeyStoreException, NoSuchAlgorithmException, IOException, CertificateException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        origKeyStore.store(baos, resultPass);

        KeyStore res = KeyStore.getInstance("jks");
        ByteArrayInputStream bois = new ByteArrayInputStream(baos.toByteArray());
        res.load(bois, resultPass);
        return res;
    }
}
