package org.apache.servicemix.jbi.security.keystore.impl;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.servicemix.jbi.security.keystore.KeystoreInstance;
import org.apache.servicemix.jbi.security.keystore.KeystoreIsLocked;
import org.springframework.core.io.Resource;

/**
 * 
 * @org.apache.xbean.XBean element="keystore"
 */
public class FileKeystoreInstance implements KeystoreInstance {

    private static final Log log = LogFactory.getLog(FileKeystoreInstance.class);
    private final static String JKS = "JKS";
    
    private Resource path;
    private String name;
    private String keystorePassword;
    private Map keyPasswords = new HashMap();
    private File keystoreFile; // Only valid after startup
    
    // The following variables are the state of the keystore, which should be chucked if the file on disk changes
    private List privateKeys = new ArrayList();
    private List trustCerts = new ArrayList();
    private KeyStore keystore;
    private long keystoreReadDate = Long.MIN_VALUE;
    
    /**
     * @param keyPasswords the keyPasswords to set
     */
    public void setKeyPasswords(String keyPasswords) {
        if (keyPasswords != null) {
            String[] keys = keyPasswords.split("\\]\\!\\[");
            for (int i = 0; i < keys.length; i++) {
                String key = keys[i];
                int pos = key.indexOf('=');
                this.keyPasswords.put(key.substring(0, pos), key.substring(pos+1).toCharArray());
            }
        }
    }

    /**
     * @return the keystoreName
     */
    public String getName() {
        return name;
    }

    /**
     * @param keystoreName the keystoreName to set
     */
    public void setName(String keystoreName) {
        this.name = keystoreName;
    }

    /**
     * @param keystorePassword the keystorePassword to set
     */
    public void setKeystorePassword(String keystorePassword) {
        this.keystorePassword = keystorePassword;
    }

    /**
     * @return the keystorePath
     */
    public Resource getPath() {
        return path;
    }

    /**
     * @param keystorePath the keystorePath to set
     */
    public void setPath(Resource keystorePath) throws IOException {
        this.path = keystorePath;
        this.keystoreFile = keystorePath.getFile();
    }

    public Certificate getCertificate(String alias) {
        if (!loadKeystoreData()) {
            return null;
        }
        try {
            return keystore.getCertificate(alias);
        } catch (KeyStoreException e) {
            log.error("Unable to read certificate from keystore", e);
        }
        return null;
    }

    public String getCertificateAlias(Certificate cert) {
        if (!loadKeystoreData()) {
            return null;
        }
        try {
            return keystore.getCertificateAlias(cert);
        } catch (KeyStoreException e) {
            log.error("Unable to read retrieve alias for given certificate from keystore", e);
        }
        return null;
    }

    public Certificate[] getCertificateChain(String alias) {
        if (!loadKeystoreData()) {
            return null;
        }
        try {
            return keystore.getCertificateChain(alias);
        } catch (KeyStoreException e) {
            log.error("Unable to read certificate chain from keystore", e);
        }
        return null;
    }

    public KeyManager[] getKeyManager(String algorithm, String keyAlias) throws KeystoreIsLocked, NoSuchAlgorithmException, KeyStoreException, UnrecoverableKeyException {
        if(isKeystoreLocked()) {
            throw new KeystoreIsLocked("Keystore '"+name+"' is locked; please unlock it in the console.");
        }
        if(keystore == null || keystoreReadDate < keystoreFile.lastModified()) {
            loadKeystoreData();
        }
        KeyManagerFactory keyFactory = KeyManagerFactory.getInstance(algorithm);
        keyFactory.init(keystore, (char[]) keyPasswords.get(keyAlias));
        return keyFactory.getKeyManagers();
    }

    public PrivateKey getPrivateKey(String alias) {
        if (!loadKeystoreData()) {
            return null;
        }
        try {
            if (isKeyLocked(alias)) {
                return null;
            }
            Key key = keystore.getKey(alias, (char[]) keyPasswords.get(alias));
            if (key instanceof PrivateKey) {
                return (PrivateKey) key;
            }
        } catch (KeyStoreException e) {
            log.error("Unable to read private key from keystore", e);
        } catch (NoSuchAlgorithmException e) {
            log.error("Unable to read private key from keystore", e);
        } catch (UnrecoverableKeyException e) {
            log.error("Unable to read private key from keystore", e);
        }
        return null;
    }

    public TrustManager[] getTrustManager(String algorithm) throws KeyStoreException, NoSuchAlgorithmException, KeystoreIsLocked {
        if(isKeystoreLocked()) {
            throw new KeystoreIsLocked("Keystore '"+name+"' is locked; please unlock it in the console.");
        }
        if (!loadKeystoreData()) {
            return null;
        }
        TrustManagerFactory trustFactory = TrustManagerFactory.getInstance(algorithm);
        trustFactory.init(keystore);
        return trustFactory.getTrustManagers();
    }

    public boolean isKeyLocked(String keyAlias) {
        return keyPasswords.get(keyAlias) == null;
    }

    public boolean isKeystoreLocked() {
        return keystorePassword == null;
    }

    public String[] listPrivateKeys() {
        if (!loadKeystoreData()) {
            return null;
        }
        return (String[]) privateKeys.toArray(new String[privateKeys.size()]);
    }

    public String[] listTrustCertificates() {
        if (!loadKeystoreData()) {
            return null;
        }
        return (String[]) trustCerts.toArray(new String[trustCerts.size()]);
    }

    // ==================== Internals =====================

    private boolean loadKeystoreData() {
        if (keystoreFile == null) {
            throw new IllegalArgumentException("keystorePath not set");
        }
        if (keystoreReadDate >= keystoreFile.lastModified()) {
            return true;
        }
        if (!keystoreFile.exists() || !keystoreFile.canRead()) {
            throw new IllegalArgumentException("Invalid keystore file (" + path + " = " + keystoreFile.getAbsolutePath() + ")");
        }
        try {
            keystoreReadDate = System.currentTimeMillis();
            privateKeys.clear();
            trustCerts.clear();
            if(keystore == null) {
                keystore = KeyStore.getInstance(JKS);
            }
            InputStream in = new BufferedInputStream(new FileInputStream(keystoreFile));
            keystore.load(in, keystorePassword == null ? new char[0] : keystorePassword.toCharArray());
            in.close();
            Enumeration aliases = keystore.aliases();
            while (aliases.hasMoreElements()) {
                String alias = (String) aliases.nextElement();
                if (keystore.isKeyEntry(alias)) {
                    privateKeys.add(alias);
                } else if (keystore.isCertificateEntry(alias)) {
                    trustCerts.add(alias);
                }
            }
            return true;
        } catch (KeyStoreException e) {
            log.error("Unable to open keystore with provided password", e);
        } catch (IOException e) {
            log.error("Unable to open keystore with provided password", e);
        } catch (NoSuchAlgorithmException e) {
            log.error("Unable to open keystore with provided password", e);
        } catch (CertificateException e) {
            log.error("Unable to open keystore with provided password", e);
        }
        return false;
    }

}
