package net.sourceforge.jcetaglib.tools;

import net.sourceforge.jcetaglib.lib.CertTools;
import org.bouncycastle.asn1.DERBMPString;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.jce.interfaces.PKCS12BagAttributeCarrier;

import javax.crypto.Cipher;
import javax.crypto.EncryptedPrivateKeyInfo;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.PBEParameterSpec;
import java.io.ByteArrayInputStream;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.spec.PKCS8EncodedKeySpec;
import java.util.ArrayList;

/**
 * Tools to handle common key and keystore operations.
 *
 * @version $Id: KeyTools.java,v 1.4 2004/04/15 07:28:36 hamgert Exp $
 */
public class KeyTools {

    //private static Category cat = Category.getInstance(KeyTools.class.getName());


    private static byte[] salt = {
        (byte) 0x23, (byte) 0xc8, (byte) 0x99, (byte) 0x8c,
        (byte) 0xc4, (byte) 0xff, (byte) 0xee, (byte) 0x7d
    }; // Salt

    private static int count = 100; // Iteration count
    private static String alg = "1.2.840.113549.1.12.1.3"; // 3 key triple DES with SHA-1


    /** Prevent from creating new KeyTools object*/
    private KeyTools() {
    }

    /**
     * Creates PKCS12-file that can be imported in IE or Netscape.
     * The alias for the private key is set to 'privateKey' and the private key password is null.
     * @param alias the alias used for the key entry
     * @param privKey RSA private key
     * @param cert user certificate
     * @param cacert CA-certificate or null if only one cert in chain, in that case use 'cert'.
     * @return byte[] containing PKCS12-file in binary format
     * @exception Exception if input parameters are not OK or certificate generation fails
     */
    static public KeyStore createP12(String alias, PrivateKey privKey, X509Certificate cert, X509Certificate cacert)
            throws Exception {
        Certificate[] chain;
        if (cacert == null)
            chain = null;
        else {
            chain = new Certificate[1];
            chain[0] = cacert;
        }
        return createP12(alias, privKey, cert, chain);
    } // createP12

    /**
     * Creates PKCS12-file that can be imported in IE or Netscape.
     * The alias for the private key is set to 'privateKey' and the private key password is null.
     * @param alias the alias used for the key entry
     * @param privKey RSA private key
     * @param cert user certificate
     * @param cachain CA-certificate chain or null if only one cert in chain, in that case use 'cert'.
     * @return byte[] containing PKCS12-file in binary format
     * @exception Exception if input parameters are not OK or certificate generation fails
     */
    static public KeyStore createP12(String alias, PrivateKey privKey, X509Certificate cert, Certificate[] cachain)
            throws Exception {
        // Certificate chain, only max two levels deep unforturnately, this is a TODO:
        if (cert == null)
            throw new IllegalArgumentException("Parameter cert cannot be null.");
        int len = 1;
        if (cachain != null)
            len += cachain.length;
        Certificate[] chain = new Certificate[len];
        // To not get a ClassCastException we need to genereate a real new certificate with BC
        CertificateFactory cf = CertificateFactory.getInstance("X.509", "BC");
        chain[0] = cf.generateCertificate(new ByteArrayInputStream(cert.getEncoded()));
        if (cachain != null)
            for (int i = 0; i < cachain.length; i++) {
                X509Certificate tmpcert = (X509Certificate) cf.generateCertificate(new ByteArrayInputStream(cachain[i].getEncoded()));
                chain[i + 1] = tmpcert;
            }


        if (chain.length > 1) {
            for (int i = 1; i < chain.length; i++) {
                X509Certificate cacert = (X509Certificate) cf.generateCertificate(new ByteArrayInputStream(chain[i].getEncoded()));
                // Set attributes on CA-cert
                PKCS12BagAttributeCarrier caBagAttr = (PKCS12BagAttributeCarrier) chain[i];
                String cafriendly = CertTools.getPartFromDN(cacert.getSubjectDN().toString(), "CN");
                caBagAttr.setBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_friendlyName, new DERBMPString(cafriendly));
            }
        }
        // Set attributes on user-cert
        PKCS12BagAttributeCarrier certBagAttr = (PKCS12BagAttributeCarrier) chain[0];
        certBagAttr.setBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_friendlyName, new DERBMPString(alias));
        // in this case we just set the local key id to that of the public key
        certBagAttr.setBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_localKeyId, CertTools.createSubjectKeyId(chain[0].getPublicKey()));

        // "Clean" private key, i.e. remove any old attributes
        KeyFactory keyfact = KeyFactory.getInstance(privKey.getAlgorithm(), "BC");
        PrivateKey pk = keyfact.generatePrivate(new PKCS8EncodedKeySpec(privKey.getEncoded()));
        // Set attributes for private key
        PKCS12BagAttributeCarrier keyBagAttr = (PKCS12BagAttributeCarrier) pk;
        // in this case we just set the local key id to that of the public key
        keyBagAttr.setBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_friendlyName, new DERBMPString(alias));
        keyBagAttr.setBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_localKeyId, CertTools.createSubjectKeyId(chain[0].getPublicKey()));

        // store the key and the certificate chain
        KeyStore store = KeyStore.getInstance("PKCS12", "BC");
        store.load(null, null);
        store.setKeyEntry(alias, pk, null, chain);
        return store;
    } // createP12

    /** Retrieves the certificate chain from a keystore.
     * @param keyStore the keystore, which has been loaded and opened.
     * @param privateKeyAlias the alias of the privatekey for which the certchain belongs.
     * @return array of Certificate, length of array is 0 if no certificates are found.
     */
    public static Certificate[] getCertChain(KeyStore keyStore, String privateKeyAlias) throws KeyStoreException {

        Certificate[] certchain = keyStore.getCertificateChain(privateKeyAlias);

        if (certchain.length < 1) {
            System.out.println("Cannot load certificate chain with alias '" + privateKeyAlias + "' from keystore.");
            return certchain;
        } else if (certchain.length > 0) {
            if (CertTools.isSelfSigned((X509Certificate) certchain[certchain.length - 1])) {
                return certchain;
            }
        }

        // If we came here, we have a cert which is not root cert in 'cert'
        ArrayList array = new ArrayList();
        for (int i = 0; i < certchain.length; i++) {
            array.add(certchain[i]);
        }

        boolean stop = false;
        while (!stop) {
            X509Certificate cert = (X509Certificate) array.get(array.size() - 1);
            String ialias = CertTools.getPartFromDN(cert.getIssuerDN().toString(), "CN");
            Certificate[] chain1 = keyStore.getCertificateChain(ialias);
            if (chain1 == null) {
                stop = true;
            } else {
                if (chain1.length == 0) {
                    System.out.println("No RootCA certificate found!");
                    stop = true;
                }
                for (int j = 0; j < chain1.length; j++) {
                    array.add(chain1[j]);
                    // If one cert is slefsigned, we have found a root certificate, we don't need to go on anymore
                    if (CertTools.isSelfSigned((X509Certificate) chain1[j]))
                        stop = true;
                }
            }
        }
        Certificate[] ret = new Certificate[array.size()];
        for (int i = 0; i < ret.length; i++) {
            ret[i] = (X509Certificate) array.get(i);
        }
        return ret;
    } // getCertChain

    /** decrypts an RSA private key.
     * @param wrappedKey the key in bytes
     * @param password password as string
     * @return PrivateKey structure
     */
    public static PrivateKey decryptPrivateKey(byte[] wrappedKey, String password)
            throws Exception {
        PBEParameterSpec defParams = new PBEParameterSpec(salt, count);

        AlgorithmParameters params = AlgorithmParameters.getInstance(alg, "BC");

        params.init(defParams);

        //
        // set up the key
        //

        EncryptedPrivateKeyInfo privKeyInfo = new EncryptedPrivateKeyInfo(params, wrappedKey);


        PBEKeySpec pbeSpec = new PBEKeySpec(password.toCharArray());
        SecretKeyFactory keyFact = SecretKeyFactory.getInstance(alg, "BC");
        Cipher cipher = Cipher.getInstance(alg, "BC");

        cipher.init(Cipher.DECRYPT_MODE, keyFact.generateSecret(pbeSpec), privKeyInfo.getAlgParameters());


        PKCS8EncodedKeySpec keySpec = privKeyInfo.getKeySpec(cipher);

        KeyFactory kf = KeyFactory.getInstance("RSA");
        PrivateKey privKey = kf.generatePrivate(keySpec);

        return privKey;
    } //decryptPrivateKey

} // KeyTools//