/*
 * Class: ${CLASS}
 *
 * Author: Alex Usun
 *
 * Date: 31/7/2005
 */
package com.jscape.util.keystore;

import com.jscape.util.StringTools;
import com.jscape.util.collections.EnumerationIterator;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.file.Files;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchProviderException;
import java.security.cert.Certificate;
import java.util.LinkedList;
import java.util.List;
import java.util.UUID;

import static com.jscape.util.collections.Algorithms.toArray;

/**
 * @author Alex Usun
 * @version 1.0
 */
public class KeyStoreTools {

    public static KeyStore create(String type, String provider)
            throws KeyStoreException, NoSuchProviderException {
        return StringTools.isValidString(provider)
               ? KeyStore.getInstance(type, provider)
               : KeyStore.getInstance(type);
    }

    public static KeyStore create(String type, String password, String provider)
            throws OperationException {
        try {
            KeyStore keyStore = create(type, provider);
            keyStore.load(null, password.toCharArray());

            return keyStore;
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static KeyStore load(InputStream in, String password, String type, String provider)
            throws OperationException {
        try {
            KeyStore keyStore = create(type, provider);
            keyStore.load(in, password.toCharArray());

            return keyStore;
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static KeyStore load(File file, String password, String type, String provider)
            throws OperationException {
        try (InputStream in = Files.newInputStream(file.toPath())) {
            return load(in, password, type, provider);
        } catch (IOException e) {
            throw new OperationException(e);
        }
    }

    public static KeyStore load(File file, String password)
            throws OperationException {
        return load(file, password, "JKS", null);
    }

    public static void store(KeyStore keyStore, File file, String password)
            throws OperationException {
        try (OutputStream out = Files.newOutputStream(file.toPath())) {
            keyStore.store(out, password.toCharArray());
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static KeyStore copyUnprotected(KeyStore keyStore, String password)
            throws OperationException {
        try {
            KeyStore destination = create(keyStore.getType(), "", keyStore.getProvider().getName());

            for (String alias : certificateAliasesOf(keyStore)) {
                Certificate certificate = keyStore.getCertificate(alias);
                destination.setCertificateEntry(alias, certificate);
            }

            for (String alias : keyAliasesOf(keyStore)) {
                Key key = keyStore.getKey(alias, password.toCharArray());
                Certificate[] certificates = keyStore.getCertificateChain(alias);
                destination.setKeyEntry(alias, key, new char[0], certificates);
            }

            return destination;
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static String[] aliasesOf(KeyStore keyStore)
            throws OperationException {
        try {
            return toArray(EnumerationIterator.iterator(keyStore.aliases()), String.class);
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static String[] keyAliasesOf(KeyStore keyStore)
            throws OperationException {
        try {
            List<String> aliases = new LinkedList<>();

            for (String alias : aliasesOf(keyStore)) {
                if (keyStore.isKeyEntry(alias)) {
                    aliases.add(alias);
                }
            }

            return aliases.toArray(new String[0]);
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static String[] certificateAliasesOf(KeyStore keyStore)
            throws OperationException {
        try {
            List<String> aliases = new LinkedList<>();

            for (String alias : aliasesOf(keyStore)) {
                if (keyStore.isCertificateEntry(alias)) {
                    aliases.add(alias);
                }
            }

            return aliases.toArray(new String[0]);
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static List<Certificate> certificatesOf(KeyStore keyStore)
            throws OperationException {
        try {
            List<Certificate> certificates = new LinkedList<>();

            for (String alias : aliasesOf(keyStore)) {
                if (keyStore.isCertificateEntry(alias)) {
                    Certificate certificate = keyStore.getCertificate(alias);
                    certificates.add(certificate);
                }
            }

            return certificates;
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static boolean hasKeys(KeyStore keyStore)
            throws OperationException {
        return keyAliasesOf(keyStore).length > 0;
    }

    public static boolean hasCertificates(KeyStore keyStore)
            throws OperationException {
        return certificateAliasesOf(keyStore).length > 0;
    }

    public static KeyStore put(KeyStore keyStore, Certificate... certificates)
            throws OperationException {
        if (certificates == null) {
            return keyStore;
        }

        try {
            for (Certificate certificate : certificates) {
                String alias = UUID.randomUUID().toString().replace("-", "");
                keyStore.setCertificateEntry(alias, certificate);
            }

            return keyStore;
        } catch (KeyStoreException e) {
            throw new OperationException(e);
        }
    }

    private KeyStoreTools() {
        // empty
    }

    public static class OperationException
            extends Exception {

        public OperationException() {
        }

        public OperationException(String message) {
            super(message);
        }

        public OperationException(String message, Throwable cause) {
            super(message, cause);
        }

        public OperationException(Throwable cause) {
            super(cause.getMessage() != null ? cause.getMessage() : cause.toString(), cause);
        }
    }

}
