/*
 * Copyright (c) 2005, 2012, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * 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 General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

import java.io.*;
import java.security.Key;
import java.security.KeyException;
import java.security.PublicKey;
import java.security.cert.*;
import java.util.*;
import javax.crypto.SecretKey;
import javax.xml.crypto.*;
import javax.xml.crypto.dsig.*;
import javax.xml.crypto.dom.*;
import javax.xml.crypto.dsig.keyinfo.*;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.DocumentBuilder;
import org.w3c.dom.Document;
import org.w3c.dom.Node;
import org.w3c.dom.Element;
import org.w3c.dom.traversal.*;
import sun.security.util.DerValue;
import sun.security.x509.X500Name;

/**
 * This is a class which supplies several KeySelector implementations
 */
class KeySelectors {

    /**
     * KeySelector which would always return the secret key specified in its
     * constructor.
     */
    static class SecretKeySelector extends KeySelector {
        private SecretKey key;
        SecretKeySelector(byte[] bytes) {
            key = wrapBytes(bytes);
        }
        SecretKeySelector(SecretKey key) {
            this.key = key;
        }

        public KeySelectorResult select(KeyInfo ki,
                                        KeySelector.Purpose purpose,
                                        AlgorithmMethod method,
                                        XMLCryptoContext context)
            throws KeySelectorException {
            return new SimpleKSResult(key);
        }

        private SecretKey wrapBytes(final byte[] bytes) {
            return new SecretKey() {
                public String getFormat() {
                    return "RAW";
                }

                public String getAlgorithm() {
                    return "Secret key";
                }

                public byte[] getEncoded() {
                    return bytes.clone();
                }
            };
        }
    }

    /**
     * KeySelector which would retrieve the X509Certificate out of the
     * KeyInfo element and return the public key.
     * NOTE: If there is an X509CRL in the KeyInfo element, then revoked
     * certificate will be ignored.
     */
    static class RawX509KeySelector extends KeySelector {

        public KeySelectorResult select(KeyInfo keyInfo,
                                        KeySelector.Purpose purpose,
                                        AlgorithmMethod method,
                                        XMLCryptoContext context)
            throws KeySelectorException {
            if (keyInfo == null) {
                throw new KeySelectorException("Null KeyInfo object!");
            }
            // search for X509Data in keyinfo
            Iterator iter = keyInfo.getContent().iterator();
            while (iter.hasNext()) {
                XMLStructure kiType = (XMLStructure) iter.next();
                if (kiType instanceof X509Data) {
                    X509Data xd = (X509Data) kiType;
                    Object[] entries = xd.getContent().toArray();
                    X509CRL crl = null;
                    // Looking for CRL before finding certificates
                    for (int i = 0; (i<entries.length&&crl != null); i++) {
                        if (entries[i] instanceof X509CRL) {
                            crl = (X509CRL) entries[i];
                        }
                    }
                    Iterator xi = xd.getContent().iterator();
                    boolean hasCRL = false;
                    while (xi.hasNext()) {
                        Object o = xi.next();
                        // skip non-X509Certificate entries
                        if (o instanceof X509Certificate) {
                            if ((purpose != KeySelector.Purpose.VERIFY) &&
                                (crl != null) &&
                                crl.isRevoked((X509Certificate)o)) {
                                continue;
                            } else {
                                return new SimpleKSResult
                                    (((X509Certificate)o).getPublicKey());
                            }
                        }
                    }
                }
            }
            throw new KeySelectorException("No X509Certificate found!");
        }
    }

    /**
     * KeySelector which would retrieve the public key out of the
     * KeyValue element and return it.
     * NOTE: If the key algorithm doesn't match signature algorithm,
     * then the public key will be ignored.
     */
    static class KeyValueKeySelector extends KeySelector {
        public KeySelectorResult select(KeyInfo keyInfo,
                                        KeySelector.Purpose purpose,
                                        AlgorithmMethod method,
                                        XMLCryptoContext context)
            throws KeySelectorException {
            if (keyInfo == null) {
                throw new KeySelectorException("Null KeyInfo object!");
            }
            SignatureMethod sm = (SignatureMethod) method;
            List list = keyInfo.getContent();

            for (int i = 0; i < list.size(); i++) {
                XMLStructure xmlStructure = (XMLStructure) list.get(i);
                if (xmlStructure instanceof KeyValue) {
                    PublicKey pk = null;
                    try {
                        pk = ((KeyValue)xmlStructure).getPublicKey();
                    } catch (KeyException ke) {
                        throw new KeySelectorException(ke);
                    }
                    // make sure algorithm is compatible with method
                    if (algEquals(sm.getAlgorithm(), pk.getAlgorithm())) {
                        return new SimpleKSResult(pk);
                    }
                }
            }
            throw new KeySelectorException("No KeyValue element found!");
        }

        //@@@FIXME: this should also work for key types other than DSA/RSA
        static boolean algEquals(String algURI, String algName) {
            if (algName.equalsIgnoreCase("DSA") &&
                algURI.equals(SignatureMethod.DSA_SHA1)) {
                return true;
            } else if (algName.equalsIgnoreCase("RSA") &&
                (algURI.equals(SignatureMethod.RSA_SHA1) ||
                 algURI.equals
                    ("http://www.w3.org/2001/04/xmldsig-more#rsa-sha256") ||
                 algURI.equals
                    ("http://www.w3.org/2001/04/xmldsig-more#rsa-sha384") ||
                 algURI.equals
                    ("http://www.w3.org/2001/04/xmldsig-more#rsa-sha512"))) {
                return true;
            } else {
                return false;
            }
        }
    }

    /**
     * KeySelector which would perform special lookup as documented
     * by the ie/baltimore/merlin-examples testcases and return the
     * matching public key.
     */
    static class CollectionKeySelector extends KeySelector {
        private CertificateFactory cf;
        private File certDir;
        private Vector<X509Certificate> certs;
        private static final int MATCH_SUBJECT = 0;
        private static final int MATCH_ISSUER = 1;
        private static final int MATCH_SERIAL = 2;
        private static final int MATCH_SUBJECT_KEY_ID = 3;
        private static final int MATCH_CERTIFICATE = 4;

        CollectionKeySelector(File dir) {
            certDir = dir;
            try {
                cf = CertificateFactory.getInstance("X509");
            } catch (CertificateException ex) {
                // not going to happen
            }
            certs = new Vector<X509Certificate>();
            File[] files = new File(certDir, "certs").listFiles();
            for (int i = 0; i < files.length; i++) {
                try (FileInputStream fis = new FileInputStream(files[i])) {
                    certs.add((X509Certificate)cf.generateCertificate(fis));
                } catch (Exception ex) { }
            }
        }

        Vector<X509Certificate> match(int matchType, Object value,
                                      Vector<X509Certificate> pool) {
            Vector<X509Certificate> matchResult = new Vector<>();
            for (int j=0; j < pool.size(); j++) {
                X509Certificate c = pool.get(j);
                switch (matchType) {
                case MATCH_SUBJECT:
                    try {
                        if (c.getSubjectDN().equals(new X500Name((String)value))) {
                            matchResult.add(c);
                        }
                    } catch (IOException ioe) { }
                    break;
                case MATCH_ISSUER:
                    try {
                        if (c.getIssuerDN().equals(new X500Name((String)value))) {
                            matchResult.add(c);
                        }
                    } catch (IOException ioe) { }
                    break;
                case MATCH_SERIAL:
                    if (c.getSerialNumber().equals(value)) {
                        matchResult.add(c);
                    }

                    break;
                case MATCH_SUBJECT_KEY_ID:
                    byte[] extension = c.getExtensionValue("2.5.29.14");
                    if (extension != null) {
                        try {
                            DerValue derValue = new DerValue(extension);
                            DerValue derValue2 = new DerValue(derValue.getOctetString());
                            byte[] extVal = derValue2.getOctetString();

                            if (Arrays.equals(extVal, (byte[]) value)) {
                                matchResult.add(c);
                            }
                        } catch (IOException ex) { }
                    }
                    break;
                case MATCH_CERTIFICATE:
                    if (c.equals(value)) {
                        matchResult.add(c);
                    }
                    break;
                }
            }
            return matchResult;
        }

        public KeySelectorResult select(KeyInfo keyInfo,
                                        KeySelector.Purpose purpose,
                                        AlgorithmMethod method,
                                        XMLCryptoContext context)
            throws KeySelectorException {
            if (keyInfo == null) {
                throw new KeySelectorException("Null KeyInfo object!");
            }
            Iterator iter = keyInfo.getContent().iterator();
            while (iter.hasNext()) {
                XMLStructure xmlStructure = (XMLStructure) iter.next();
                try {
                    if (xmlStructure instanceof KeyName) {
                        String name = ((KeyName)xmlStructure).getName();
                        PublicKey pk = null;
                        File certFile = new File(new File(certDir, "certs"),
                                                 name.toLowerCase() + ".crt");
                        try (FileInputStream fis = new FileInputStream(certFile)) {
                            // Lookup the public key using the key name 'Xxx',
                            // i.e. the public key is in "certs/xxx.crt".
                            X509Certificate cert = (X509Certificate)
                                cf.generateCertificate(fis);
                            pk = cert.getPublicKey();
                        } catch (FileNotFoundException e) {
                            // assume KeyName contains subject DN and search
                            // collection of certs for match
                            Vector<X509Certificate> result =
                                match(MATCH_SUBJECT, name, certs);
                            int numOfMatches = (result==null? 0:result.size());
                            if (numOfMatches != 1) {
                                throw new KeySelectorException
                                    ((numOfMatches==0?"No":"More than one") +
                                     " match found");
                            }
                            pk = result.get(0).getPublicKey();
                        }
                        return new SimpleKSResult(pk);
                    } else if (xmlStructure instanceof RetrievalMethod) {
                        // Lookup the public key using the retrievel method.
                        // NOTE: only X509Certificate type is supported.
                        RetrievalMethod rm = (RetrievalMethod) xmlStructure;
                        String type = rm.getType();
                        if (type.equals(X509Data.RAW_X509_CERTIFICATE_TYPE)) {
                            String uri = rm.getURI();
                            try (FileInputStream fis =
                                 new FileInputStream(new File(certDir, uri))) {
                                X509Certificate cert = (X509Certificate)
                                    cf.generateCertificate(fis);
                                return new SimpleKSResult(cert.getPublicKey());
                            }
                        } else {
                            throw new KeySelectorException
                                ("Unsupported RetrievalMethod type");
                        }
                    } else if (xmlStructure instanceof X509Data) {
                        List content = ((X509Data)xmlStructure).getContent();
                        int size = content.size();
                        Vector<X509Certificate> result = null;
                        // Lookup the public key using the information
                        // specified in X509Data element, i.e. searching
                        // over the collection of certificate files under
                        // "certs" subdirectory and return those match.
                        for (int k = 0; k<size; k++) {
                            Object obj = content.get(k);
                            if (obj instanceof String) {
                                result = match(MATCH_SUBJECT, obj, certs);
                            } else if (obj instanceof byte[]) {
                                result = match(MATCH_SUBJECT_KEY_ID, obj,
                                               certs);
                            } else if (obj instanceof X509Certificate) {
                                result = match(MATCH_CERTIFICATE, obj, certs);
                            } else if (obj instanceof X509IssuerSerial) {
                                X509IssuerSerial is = (X509IssuerSerial) obj;
                                result = match(MATCH_SERIAL,
                                               is.getSerialNumber(), certs);
                                result = match(MATCH_ISSUER,
                                               is.getIssuerName(), result);
                            } else {
                                throw new KeySelectorException("Unsupported X509Data: " + obj);
                            }
                        }
                        int numOfMatches = (result==null? 0:result.size());
                        if (numOfMatches != 1) {
                            throw new KeySelectorException
                                ((numOfMatches==0?"No":"More than one") +
                                 " match found");
                        }
                        return new SimpleKSResult(result.get(0).getPublicKey());
                    }
                } catch (Exception ex) {
                    throw new KeySelectorException(ex);
                }
            }
            throw new KeySelectorException("No matching key found!");
        }
    }

    static class ByteUtil {

        private static String mapping = "0123456789ABCDEF";
        private static int numBytesPerRow = 6;

        private static String getHex(byte value) {
            int low = value & 0x0f;
            int high = ((value >> 4) & 0x0f);
            char[] res = new char[2];
            res[0] = mapping.charAt(high);
            res[1] = mapping.charAt(low);
            return new String(res);
        }

        static String dumpArray(byte[] in) {
            int numDumped = 0;
            StringBuffer buf = new StringBuffer(512);
            buf.append("{");
            for (int i=0;i<(in.length/numBytesPerRow); i++) {
                for (int j=0; j<(numBytesPerRow); j++) {
                    buf.append("(byte)0x" + getHex(in[i*numBytesPerRow+j]) +
                               ", ");
                }
                numDumped += numBytesPerRow;
            }
            while (numDumped < in.length) {
                buf.append("(byte)0x" + getHex(in[numDumped]) + " ");
                numDumped += 1;
            }
            buf.append("}");
            return buf.toString();
        }
    }
}

class SimpleKSResult implements KeySelectorResult {
    private final Key key;

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

    public Key getKey() { return key; }
}