package dk.cryptography;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.Principal;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.SignatureException;
import java.security.cert.Certificate;
import java.security.cert.CertificateEncodingException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Date;
import java.util.Enumeration;
import java.util.Vector;

/**
 * About: general management of PKCS#12 certificate format.<br><br>
 * 
 * @author cryptography.dk (http://www.cryptography.dk)
 *
 * Copyright www.cryptography.dk /2004
 * Revision 1.0, stigv@hotmail.com
 * 
 * Web reference: http://www.cryptography.dk
 *                http://www.topsecurity.dk
 *                http://www.compression.dk
 *
 * Open source software under the terms of the
 * GNU Lesser General Public License as published by the 
 * Free Software Foundation. Please observe: 
 * http://www.gnu.org/licenses/lgpl.txt.
 *
 * Permission to use, copy, modify, and distribute this software
 * for any purpose and without fee is hereby granted, provided
 * that the above copyright notices appear in all copies and that
 * both the copyright notice and this permission notice appear in
 * supporting documentation.
 *
 * This code is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Lesser General Public License for more details.
 * 
 */
public class Pkcs12CertificateManager {
    private java.security.cert.Certificate[] chain;
    private String alias;
    private File file;
    private KeyStore keyStore;
    private String[] aliases;
    private boolean[] keyusage = null;
    
    private char[] pkcs12Password = null;


    /**
     * @param file
     */
    public Pkcs12CertificateManager(File file) {
        this.file = file;
    }


    public void alertMsg(String msg) {
        System.out.println(msg);
    }

    public void errorMsg(String msg) {
        System.out.println(msg);
    }

    public void debugMsg(String msg) {
        //System.out.println(msg);
    }


    /**
     * Constructor takes a PKCS#12 formatted keystore file as argument.<br><br>
     * Initializes internal class variables 
     * 
     * @param file
     * @param password
     * @throws IOException
     * @throws Exception
     */
    public Pkcs12CertificateManager(File file, char[] keyPassword) throws IOException, Exception {
        this.file = file;
        getKeyStore(keyPassword);
        getLongestCertificateChain(keyPassword);
    }

    public File getFile() {
        return file;
    }

    public String toString() {
        return getFile().toString();
    }


    /**
     * Returns the longest certificate chain available in the keystore.<br><br> 
     * Should belong to the intended 
     * 
     * Looping through the Aliases in a Key Store - hereby returning the _first_ signable 
     * certificate alias (with an existing private key). Ignores certificates in 
     * keystore from the certificate chain. Returns null, if none exist
     * 
     * Modifies global class variables:<br> 
     *   chain<br><br>
     *  
     * @param pw
     * @return
     * @throws IOException
     * @throws Exception
     */
    private X509Certificate getCert(char[] keyPassword) throws IOException, Exception {
        if (chain == null) {
            try {
                getLongestCertificateChain(keyPassword);
            } catch (IOException iox) {
                alertMsg("Error: incorrect password");
                throw iox;
            } catch (Exception ex) {
                errorMsg("Failed to retrieve certificate information");
                ex.printStackTrace();
                throw ex;
            }
        }
        if(chain == null)
            throw new Exception("Information unavailable. No certificate chain recorded in the keystore.");
        
        return ((chain!=null && chain.length>0 && chain[0] instanceof X509Certificate) ? (X509Certificate) chain[0] : null);
    }



    public String getCertificateID() {
       return getFile().toString().replace('\\', '/');
    }
    
    /**
     * Returns the key associated with the given alias, using the given password to recover it. 
     * 
     * http://java.sun.com/j2se/1.4.2/docs/api/java/security/KeyStore.html#getKey(java.lang.String, char[])
     * 
     * @param password
     * @return
     * @throws IOException
     * @throws Exception
     */
    public PrivateKey getPrivateKey(char[] password) throws IOException, Exception {
        PrivateKey privateKey = (PrivateKey) getKeyStore(password).getKey(getKeyAlias(password), password);
        return privateKey;
    }


    /**
     * Retrieves public key for the _first_ signable certificate alias (with an existing private key).
     * Looping through the certificate chain and stopping at first occurrence.<br>
     * 
     * In practice this means the certificate for the certificate owner - the person holding and using the
     * certificate.<br><br>
     *  
     * @param password
     * @return
     * @throws IOException
     * @throws Exception
     */
    public PublicKey getPublicKey(char[] password) throws IOException, Exception {
        Certificate[] aliasCerts = getKeyStore(password).getCertificateChain(getKeyAlias(password));
        return (aliasCerts.length>0 ? aliasCerts[0].getPublicKey() : null);
    }

    /**
     * Returns validation certificate - a certificate signed by the issuer of certificates contained in
     * the keystore. This certificate is used to prove the origin validity of the other certificates.<br><br> 
     *  
     * @param password
     * @return
     * @throws IOException
     * @throws Exception
     */
    public Certificate getOwnerCertificate(char[] password) throws IOException, Exception {
        Certificate[] aliasCerts = getKeyStore(password).getCertificateChain(getKeyAlias(password));
        return (aliasCerts.length>0 ? aliasCerts[0] : null);
    }



    /**
     * Returns validation certificate - a certificate signed by the issuer of certificates contained in
     * the keystore. This certificate is used to prove the origin validity of the other certificates.<br><br> 
     *  
     * @param password
     * @return
     * @throws IOException
     * @throws Exception
     */
    public Certificate getValidaterCertificate(char[] password) throws IOException, Exception {
        Certificate[] aliasCerts = getKeyStore(password).getCertificateChain(getKeyAlias(password));
        return (aliasCerts.length>1 ? aliasCerts[aliasCerts.length-1] : null);
    }




    /**
     * Returns all aliases - both for key and certificate entries in the keystore. You are left
     * for yourself to sort out what kind it is.
     * 
     * Used twice. Thus subroutined.
     * 
     *  
     * @param keystore
     * @return
     * @throws KeyStoreException
     */
    private static String[] getAllAliases(KeyStore keystore) throws KeyStoreException {
        Enumeration enum = keystore.aliases();
        Vector list = new Vector();
        while (enum.hasMoreElements()) {
            list.addElement(enum.nextElement().toString());
        }
        String [] aliasArray = new String[list.size()];
        list.copyInto(aliasArray);
        return aliasArray;
    }


    /**
     * Looping through the Aliases in a Key Store - hereby returning the _first_ signable 
     * certificate alias (with an existing private key). Ignores certificates in 
     * keystore from the certificate chain. Returns null, if none exist  
     * 
     * Modifies global class variables:<br> 
     *   alias<br><br>
     * 
     * A key store is a collection of keys and certificates. Each key or certificate has 
     * a unique alias used to identify that key or certificate in the key store. An 
     * example to list all the aliases in a key store: 
     * http://javaalmanac.com/egs/java.security/ListAliases.html?l=rel
     * 
     */
    public synchronized String getKeyAlias(char[] password) throws IOException, Exception {
        if (alias == null) {
            KeyStore keystore = getKeyStore(password);            
            Enumeration enum = keystore.aliases(); // List the aliases
            while(enum.hasMoreElements()) {
                String keyStoreAlias = enum.nextElement().toString();
                if(keystore.isKeyEntry(keyStoreAlias)) { //Does alias refer to a private key?
                  alias=keyStoreAlias;
                  break;
        }   }   }
        return alias;
    }


    /**
     * Listing all aliases in a Key Store.
     * 
     * Modifies global class variables:<br> 
     *   aliases<br><br>
     * 
     * A key store is a collection of keys and certificates. Each key or certificate has 
     * a unique alias used to identify that key or certificate in the key store. An 
     * example to list all the aliases in a key store: 
     * http://javaalmanac.com/egs/java.security/ListAliases.html?l=rel
     * 
     * @param store
     * @return
     * @throws KeyStoreException
     */
    public synchronized String[] getAliases(char[] password) throws IOException, Exception {
        if (aliases == null) {
            KeyStore keystore = getKeyStore(password);
            aliases = getAllAliases(keyStore);
        }
        return aliases;
    }

    /**
     * Retrieving key store from a file.
     * 
     * Modifies global class variables:<br> 
     *   keystore<br><br>
     * 
     * http://javaalmanac.com/egs/java.security/GetCertFromKeyStore.html
     * 
     * @param password
     * @return
     * @throws IOException
     * @throws Exception
     */
    public synchronized KeyStore getKeyStore(char[] password) throws IOException, Exception {
        if (keyStore == null) {
            FileInputStream fin = new FileInputStream(file);
            KeyStore store = KeyStore.getInstance("PKCS12");
            try {
                store.load(fin, password);                
            } 
            finally {
                try {
                    fin.close();
                } catch (IOException e) { }
            }
            keyStore = store;
        }
        return keyStore;
    }

    /**
     * Return subject name of the certificate holder.<br><br>
     * 
     * Strictly, you don't need to enter the password each time. If the 'cert' strain has been initialized
     * when constructing the class, the password in the argument won't be used (again).<br><br> 
     *  
     * @param keyPassword
     * @return Principal
     * @throws Exception
     */
    public Principal getSubjectDN(char[] keyPassword) throws Exception {
        return getCert(keyPassword).getSubjectDN();
    }

    public Principal getIssuerDN(char[] keyPassword) throws Exception {
        return getCert(keyPassword).getIssuerDN();
    };

    public BigInteger getSerialNumber(char[] keyPassword) throws Exception {
        return getCert(keyPassword).getSerialNumber();
    }

    public Date getNotBefore(char[] keyPassword) throws Exception {
        return getCert(keyPassword).getNotBefore();
    }

    public Date getNotAfter(char[] keyPassword) throws Exception {
        return getCert(keyPassword).getNotAfter();
    }

    public int getVersion(char[] keyPassword) throws Exception {
        return getCert(keyPassword).getVersion();
    }

    public byte[] getCertificate(char[] keyPassword) throws Exception, IOException {
        return getCert(keyPassword).getEncoded();
    }


    /**
     * Creating a Signature using the owners private key in the PKCS#12 keystore.<br><br>
     * 
     * http://javaalmanac.com/egs/java.security/CreateSig.html
     * 
     * @param toBeSigned
     * @param password
     * @return
     * @throws IOException
     * @throws Exception
     */
    public byte[] sign(byte[] dataToSign, char[] keyPassword) throws IOException, Exception {
        if (canSign(keyPassword)) {
            Signature sig = Signature.getInstance("SHA1withRSA");
            sig.initSign(this.getPrivateKey(keyPassword));
            sig.update(dataToSign);
            byte[] signature = sig.sign();
            return signature;
        } else
            throw new CertificateException("Certificate not applicable to signing");
    }

    /**
     * Creating a Keyed Digest 
     * 
     * http://javaalmanac.com/egs/java.security/Digest.html
     *   
     * @param dataToDigest
     * @return
     * @throws NoSuchProviderException
     * @throws NoSuchAlgorithmException
     */
    public byte[] digest(byte[] dataToDigest) throws NoSuchProviderException, NoSuchAlgorithmException {
        MessageDigest sha = MessageDigest.getInstance("SHA1");
        return sha.digest(dataToDigest);
    }


    /**
     * Return the validation certificate used to verity the actual certificate issued in this keystore.<br><br> 
     *  
     * @param password
     * @return
     */
    public String getCertificateVerifier(char[] password) {
        String verifier = null;
        try {
            Certificate[] aliasCerts = getKeyStore(password).getCertificateChain(getKeyAlias(password));
            if(aliasCerts!=null && aliasCerts.length>0 && aliasCerts[aliasCerts.length-1] instanceof X509Certificate) {
                Principal p =(((X509Certificate)aliasCerts[aliasCerts.length-1]).getSubjectDN());
                verifier = p.toString();
            }
        }
        catch(IOException iox) {
            iox.printStackTrace();
        }
        catch(Exception ex) {
            ex.printStackTrace();
        }
        return verifier;
    }




    /**
     * Returns the _longest_ certificate chain (ordered with the user's certificate first 
     * and the root certificate authority last), or null if the given alias does not exist 
     * or does not contain a certificate chain (i.e., the given alias identifies either 
     * a trusted certificate entry or a key entry without a certificate chain).<br><br>
     * 
     * Usually the longest certificate chain is associated to the principal owner of the keystore.<br><br>
     * 
     * Modifies global class variables:<br>
     *   chain<br><br>
     * 
     *  
     * http://java.sun.com/j2se/1.4.2/docs/api/java/security/KeyStore.html#getCertificateChain(java.lang.String)
     * 
     * @param password
     * @return
     * @throws IOException
     * @throws Exception
     */
    public byte[][] getLongestCertificateChain(char[] password) throws IOException, Exception {
        KeyStore keystore = getKeyStore(password);
        String[] aliases = getAllAliases(keystore);
        Certificate[] longestChain = null;
        for (int i = 0; i < aliases.length; i++) {
            Certificate[] certs = keystore.getCertificateChain(aliases[i]);
            if (longestChain == null || //does the alias have a longer certificate chains attanched?
                  (certs != null && longestChain.length < certs.length))
               longestChain = certs; //..then save longest chain!
        }
        byte[][] bchain = new byte[longestChain.length][0];
        
        chain = new X509Certificate[longestChain.length];
        for (int i = 0; i < longestChain.length; i++) {
            bchain[i] = longestChain[i].getEncoded();
            chain[i] = longestChain[i];
        }
        return bchain;
    }

    /**
     * TRanslates the meaning of usage to text. 
     *  
     * Gets a boolean array representing bits of the KeyUsage extension, (OID = 2.5.29.15). The key usage extension defines the purpose (e.g., encipherment, signature, certificate signing) of the key contained in the certificate. The ASN.1 definition for this is: 
     * KeyUsage ::= BIT STRING {
     *    digitalSignature        (0),
     *   nonRepudiation          (1),
     *   keyEncipherment         (2),
     *   dataEncipherment        (3),
     *   keyAgreement            (4),
     *   keyCertSign             (5),
     *   cRLSign                 (6),
     *   encipherOnly            (7),
     *   decipherOnly            (8) }
     *RFC 2459 recommends that when used, this be marked as a critical extension.
     * 
     *http://java.sun.com/j2se/1.3/docs/api/java/security/cert/X509Certificate.html#getKeyUsage()
     *
     *  
     * @return
     * @throws NoSuchProviderException
     * 
     */
    public String getKeyUsage(char[] keyPassword) throws Exception {
        String[] keyUsageExplanation = {"digitalSignature",
                                        "nonRepudiation",
                                        "keyEncipherment",
                                        "dataEncipherment",
                                        "keyAgreement",
                                        "keyCertSign",
                                        "cRLSign",
                                        "encipherOnly",
                                        "decipherOnly" };
        
        X509Certificate cert = getCert(keyPassword);
        if (cert.getKeyUsage() == null)
            return "Key usage is undefined for certificate";
        StringBuffer sb = new StringBuffer();
        for(int i=0; i<keyUsageExplanation.length; i++) {
            if(cert.getKeyUsage()[i])
                sb.append((i==0 ? "" : ", ")+keyUsageExplanation[i]);
        }
        return sb.toString().trim();
    }
    

    /**
     * According to usage, returns ability to sign.. 
     * 
     * @param pasword
     * @return boolean
     */
    public boolean canSign(char[] keyPassword) throws IOException, Exception {
        final int digitalSignature = 0;
        final int nonRepudiation = 1;

        try {
            X509Certificate cert = getCert(keyPassword);
            return (cert.getKeyUsage() == null ||
                cert.getKeyUsage()[digitalSignature] || 
                cert.getKeyUsage()[nonRepudiation]);
        } catch (Exception e) {
            return false;
        }
    }


    /**
     * Inspects certificate for usage registered.<br><br>
     *  
     * Modifies global class variables:<br> 
     *   keyusage<br><br>
     * 
     * @return
     * @throws Exception
     * @throws IOException
     */
    public boolean[] getIntendedKeyUsage(char[] keyPassword) throws Exception, IOException {
        if (keyusage == null) {
            X509Certificate cert = getCert(keyPassword);
            boolean ku[];
            if ((ku = cert.getKeyUsage()) == null ) {
                return null;
            }
            keyusage = ku;            
        }
        return keyusage;
    }
    
    
    /**
     * Validates the issuer certificate from last in the key chain - with the CA certificate.
     * Throws exception if anythings fails - otherwise the verification may be assumed to be 
     * successfull.<br><br>
     * 
     * @param keyPassword
     * @param publicKeyCA
     * @throws CertificateException
     * @throws NoSuchAlgorithmException
     * @throws SignatureException
     * @throws NoSuchProviderException
     * @throws InvalidKeyException
     * @throws IOException
     * @throws Exception
     */
    public void validateCA(char[] keyPassword, PublicKey publicKeyCA) throws 
                CertificateException, 
                NoSuchAlgorithmException, 
                SignatureException, 
                NoSuchProviderException, 
                InvalidKeyException, 
                IOException,
                Exception {
        
        getValidaterCertificate(keyPassword).verify(publicKeyCA);
    }

    /**
     * Exporting certificate to text format. May be saved to a file.<br><br>
     * 
     * If the certificate is in a key store, it can exported using "keytool" 
     * (export in text format):<br><br>
     * keytool -storepass my-keystore-password -alias myalias -export -rfc -file outfilename.cer
     * 
     * @param cert
     * @return
     */    
    public static String exportCert(Certificate cert) {
        try {
            // Get the encoded form which is suitable for exporting
            byte[] buf = cert.getEncoded();
            StringBuffer sb = new StringBuffer();
            sb.append("-----BEGIN CERTIFICATE-----\n");
            sb.append(new sun.misc.BASE64Encoder().encode(buf));
            sb.append("\n-----END CERTIFICATE-----\n");
            return sb.toString().trim();
        } catch (CertificateEncodingException e) {
            return null;
        } 
    }
}