package cl.altiuz.reports;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.security.Key;
import java.security.KeyException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509CertSelector;
import java.security.cert.X509Certificate;
import java.util.Enumeration;
import java.util.NoSuchElementException;

import javax.security.auth.x500.X500Principal;
import javax.xml.crypto.AlgorithmMethod;
import javax.xml.crypto.KeySelector;
import javax.xml.crypto.KeySelectorException;
import javax.xml.crypto.KeySelectorResult;
import javax.xml.crypto.NodeSetData;
import javax.xml.crypto.OctetStreamData;
import javax.xml.crypto.XMLCryptoContext;
import javax.xml.crypto.XMLStructure;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyName;
import javax.xml.crypto.dsig.keyinfo.KeyValue;
import javax.xml.crypto.dsig.keyinfo.RetrievalMethod;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.keyinfo.X509IssuerSerial;

public class X509KeySelector extends KeySelector {

    private KeyStore keyStore;
    private boolean isPublicKey;

    public X509KeySelector(String keystorePath, String keystorePassword,
            boolean isPublicKey) throws KeyStoreException,
            NoSuchAlgorithmException, CertificateException, IOException {
        File file = new File(keystorePath);
        FileInputStream is = new FileInputStream(file);
        keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
        keyStore.load(is, keystorePassword.toCharArray());
        this.isPublicKey = isPublicKey;
    }

    public X509KeySelector(String keystorePath, String keystorePassword)
            throws KeyStoreException, NoSuchAlgorithmException,
            CertificateException, IOException {
        File file = new File(keystorePath);
        FileInputStream is = new FileInputStream(file);
        keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
        keyStore.load(is, keystorePassword.toCharArray());
        this.isPublicKey = true;
    }

    @Override
    public KeySelectorResult select(KeyInfo ki, KeySelector.Purpose p,
            AlgorithmMethod am, XMLCryptoContext ctx)
            throws KeySelectorException {

        KeySelectorResult result = null;
        SignatureMethod sm = (SignatureMethod) am;
        if (ki == null) {
            throw new IllegalArgumentException("No KeyInfo found");
        }
        for (Object obj : ki.getContent()) {
            XMLStructure xs = (XMLStructure) obj;
            if (xs instanceof X509Data) {
                try {
                    result = select((X509Data) xs, sm);
                } catch (UnrecoverableKeyException e) {
                    throw new KeySelectorException(e);
                } catch (NoSuchAlgorithmException e) {
                    throw new KeySelectorException(e);
                }
            } else if (xs instanceof KeyName) {
                KeyName kn = (KeyName) xs;
                try {
                    Certificate cert = keyStore.getCertificate(kn.getName());
                    if (cert != null
                            && equals(sm.getAlgorithm(), cert.getPublicKey()
                                    .getAlgorithm())) {
                        result = new SimpleKeySelectorResult(
                                cert.getPublicKey());
                    }
                } catch (KeyStoreException e) {
                    throw new KeySelectorException(e);
                }
            } else if (xs instanceof RetrievalMethod) {
                RetrievalMethod rm = (RetrievalMethod) xs;
                try {
                    if (rm.getType().equals(X509Data.RAW_X509_CERTIFICATE_TYPE)) {
                        OctetStreamData data = (OctetStreamData) rm
                                .dereference(ctx);
                        CertificateFactory cf = CertificateFactory
                                .getInstance(" X.509 ");
                        X509Certificate cert = (X509Certificate) cf
                                .generateCertificate(data.getOctetStream());
                        result = select(cert, sm);
                    } else if (rm.getType().equals(X509Data.TYPE)) {
                        NodeSetData nd = (NodeSetData) rm.dereference(ctx);
                        System.out.println(" DAS conversion is TBD ");
                    }
                } catch (Exception e) {
                    throw new KeySelectorException(e);
                }
            } else if (xs instanceof KeyValue) {
                PublicKey pk = null;
                try {
                    pk = ((KeyValue) xs).getPublicKey();
                } catch (KeyException ke) {
                    throw new KeySelectorException(ke);
                }
                if (equals(sm.getAlgorithm(), pk.getAlgorithm())) {
                    result = new SimpleKeySelectorResult(pk);
                }
            }
            if (result != null) {
                return result;
            }
        }
        throw new KeySelectorException(" Key not found ");
    }

    // Internal

    KeySelectorResult select(X509Data data, SignatureMethod sm)
            throws KeySelectorException, UnrecoverableKeyException,
            NoSuchAlgorithmException {
        String oid = getOid(sm.getAlgorithm());
        KeySelectorResult result = null;
        for (Object obj : data.getContent()) {
            try {
                if (obj instanceof X509Certificate) {
                    return select((X509Certificate) obj, sm);
                } else if (obj instanceof X509IssuerSerial) {
                    X509IssuerSerial xis = (X509IssuerSerial) obj;
                    X509CertSelector xcs = new X509CertSelector();
                    xcs.setSubjectPublicKeyAlgID(oid);
                    xcs.setSerialNumber(xis.getSerialNumber());
                    xcs.setIssuer(new X500Principal(xis.getIssuerName())
                            .getName());
                    return select(xcs);
                } else if (obj instanceof String) {
                    String sn = (String) obj;
                    X509CertSelector xcs = new X509CertSelector();
                    xcs.setSubjectPublicKeyAlgID(oid);
                    xcs.setSubject(new X500Principal(sn).getName());
                    return select(xcs);
                } else if (obj instanceof byte[]) {
                    byte[] ski = (byte[]) obj;
                    X509CertSelector xcs = new X509CertSelector();
                    xcs.setSubjectPublicKeyAlgID(oid);
                    byte[] encodedSki = new byte[ski.length + 2];
                    encodedSki[0] = 0x04; // OCTET STRING tag value
                    encodedSki[1] = (byte) ski.length;
                    System.arraycopy(ski, 0, encodedSki, 2, ski.length);
                    xcs.setSubjectKeyIdentifier(encodedSki);
                    return select(xcs);
                }
            } catch (IOException e) {
                throw new KeySelectorException(e);
            } catch (KeyStoreException e) {
                throw new KeySelectorException(e);
            } catch (NoSuchElementException e) {
                throw new KeySelectorException(e);
            }
        }
        return null;
    }

    /**
     * Return the certificate in the key store that matches the specified
     * selector
     * 
     * @throws NoSuchElementException
     *             if there is no match
     */
    KeySelectorResult select(X509CertSelector cs) throws KeyStoreException,
            NoSuchElementException {
        for (Enumeration e = keyStore.aliases(); e.hasMoreElements();) {
            Certificate cert = keyStore
                    .getCertificate((String) e.nextElement());
            if (cert != null && cs.match(cert)) {
                return new SimpleKeySelectorResult(cert.getPublicKey());
            }
        }
        throw new NoSuchElementException(cs.getSubject().toString());
    }

    /**
     * Return the certificate in the key store that matches the specified X509
     * certificate (or null if there is no match).
     * 
     * @throws NoSuchElementException
     *             if there is no match
     * @throws NoSuchAlgorithmException
     * @throws UnrecoverableKeyException
     */
    KeySelectorResult select(X509Certificate xc, SignatureMethod sm)
            throws KeyStoreException, NoSuchElementException,
            UnrecoverableKeyException, NoSuchAlgorithmException {
        boolean[] keyUsage = xc.getKeyUsage();
        if (keyUsage[0]) {
            String alias = keyStore.getCertificateAlias(xc);
            if (alias != null) {
                PublicKey pk = keyStore.getCertificate(alias).getPublicKey();
                Key priv = keyStore.getKey(alias, alias.toCharArray());
                if (equals(sm.getAlgorithm(), pk.getAlgorithm())) {
                    return new SimpleKeySelectorResult(pk);
                }

                if (equals(sm.getAlgorithm(), priv.getAlgorithm())) {
                    return new SimpleKeySelectorResult(priv);
                }
                if (isPublicKey) {
                    return new SimpleKeySelectorResult(pk);
                } else {
                    return new SimpleKeySelectorResult(priv);
                }
                /*
                 * System.out.println("sm:" + sm.getAlgorithm());
                 * System.out.println("pub:" + pk.getAlgorithm());
                 * System.out.println("priv:" + priv.getAlgorithm());
                 */
            }
        }
        throw new NoSuchElementException(xc.getSubjectX500Principal()
                .toString());
    }

    /**
     * Return the OID corresponding to the Algorithm URI.
     */
    String getOid(String uri) {
        if (uri.equalsIgnoreCase(SignatureMethod.DSA_SHA1)) {
            return " 1.2.840.10040.4.1 ";
        } else if (uri.equalsIgnoreCase(SignatureMethod.RSA_SHA1)) {
            return " 1.2.840.113549.1.1 ";
        } else {
            return null;
        }
    }

    /**
     * Determine whether the algorithm URI matches the algorithm name.
     */
    boolean equals(String uri, String name) {
        if (name.equalsIgnoreCase(" DSA ")
                && uri.equalsIgnoreCase(SignatureMethod.DSA_SHA1)) {
            return true;
        } else if (name.equalsIgnoreCase(" RSA ")
                && uri.equalsIgnoreCase(SignatureMethod.RSA_SHA1)) {
            return true;
        } else {
            return false;
        }
    }
}

class SimpleKeySelectorResult implements KeySelectorResult {
    private Key key;

    SimpleKeySelectorResult(Key key) {
        this.key = key;
    }

    // Implement KeySelectorResult

    public Key getKey() {
        return key;
    }
}