package org.jboss.resteasy.security.doseta;

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.PublicKey;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;

/**
 * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
 * @version $Revision: 1 $
 */
public class KeyStoreKeyRepository {
    private KeyStore keyStore;
    private String password;

    public void init(InputStream is, String password) {
        if (password != null)
            password = password.trim();
        this.password = password;
        try {
            keyStore = KeyStore.getInstance("JKS");
            keyStore.load(is, password.toCharArray());
        } catch (KeyStoreException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        } catch (CertificateException e) {
            throw new RuntimeException(e);
        }
    }

    public KeyStoreKeyRepository(final InputStream is, final String password) {
        init(is, password);
    }

    public KeyStoreKeyRepository(final String filename, final String password) throws IOException {
        InputStream is = new FileInputStream(filename);
        init(is, password);
        is.close();
    }

    public PrivateKey getPrivateKey(Object identity) {
        try {
            Key key = keyStore.getKey(identity.toString(), password.toCharArray());
            return (PrivateKey) key;
        } catch (KeyStoreException e) {
            throw new RuntimeException(e);
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        } catch (UnrecoverableKeyException e) {
            throw new RuntimeException(e);
        }
    }

    public PublicKey getPublicKey(Object identity) {
        Certificate cert = null;
        try {
            cert = keyStore.getCertificate(identity.toString());
        } catch (KeyStoreException e) {
            throw new RuntimeException(e);
        }
        if (cert == null)
            return null;
        return cert.getPublicKey();
    }
}
