com.adobe.granite.auth.saml.util.RetrievalMethodEncryptedKeyResolver Maven / Gradle / Ivy

/*************************************************************************
 *
 * ADOBE CONFIDENTIAL
 * __________________
 *
 *  Copyright 2012 Adobe Systems Incorporated
 *  All Rights Reserved.
 *
 * NOTICE:  All information contained herein is, and remains
 * the property of Adobe Systems Incorporated and its suppliers,
 * if any.  The intellectual and technical concepts contained
 * herein are proprietary to Adobe Systems Incorporated and its
 * suppliers and are protected by trade secret or copyright law.
 * Dissemination of this information or reproduction of this material
 * is strictly forbidden unless prior written permission is obtained
 * from Adobe Systems Incorporated.
 **************************************************************************/
package com.adobe.granite.auth.saml.util;

import org.apache.xml.security.encryption.EncryptedKey;
import org.apache.xml.security.encryption.XMLCipher;
import org.apache.xml.security.encryption.XMLEncryptionException;
import org.apache.xml.security.keys.keyresolver.KeyResolverException;
import org.apache.xml.security.keys.keyresolver.KeyResolverSpi;
import org.apache.xml.security.keys.storage.StorageResolver;
import org.w3c.dom.Element;

import javax.crypto.SecretKey;
import javax.xml.namespace.NamespaceContext;
import javax.xml.namespace.QName;
import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathConstants;
import javax.xml.xpath.XPathExpressionException;
import javax.xml.xpath.XPathFactory;
import javax.xml.xpath.XPathVariableResolver;
import java.security.Key;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class RetrievalMethodEncryptedKeyResolver extends KeyResolverSpi {
    private static org.slf4j.Logger log =
            org.slf4j.LoggerFactory.getLogger(RetrievalMethodEncryptedKeyResolver.class);

    private static final String ATTR_RETRIEVAL_METHOD = "ds:RetrievalMethod";
    private static final String TYPE_ENCRYPTED_KEY = "http://www.w3.org/2001/04/xmlenc#EncryptedKey";
    private static final String NS_PREFIX_XENC = "xenc";
    private static final String NS_URI_XENC = "http://www.w3.org/2001/04/xmlenc#";
    private final String algorithm;

    private Key decryptionKey;

    public RetrievalMethodEncryptedKeyResolver(final String algorithm, final Key decryptionKey) {
        this.algorithm = algorithm;
        this.decryptionKey = decryptionKey;
    }

    @Override
    public boolean engineCanResolve(Element element, String baseURI, StorageResolver storage) {
        return ATTR_RETRIEVAL_METHOD.equals(element.getLocalName()) && TYPE_ENCRYPTED_KEY.equals(element.getAttributeNS(null, "Type"));
    }

    @Override
    public SecretKey engineLookupAndResolveSecretKey(Element element, String baseURI, StorageResolver storage) throws KeyResolverException {
        log.debug("Lookup and resolve secret key: ", element, baseURI);
        String refURI = element.getAttributeNS(null, "URI");
        if (refURI != null && !refURI.isEmpty()) {
            XPath xpath = XPathFactory.newInstance().newXPath();
            xpath.setNamespaceContext(new NamespaceContext() {
                HashMap namespaceMap = new HashMap() {{
                    put(NS_PREFIX_XENC, NS_URI_XENC);
                }};
                HashMap prefixMap = new HashMap() {{
                    put(NS_URI_XENC, NS_PREFIX_XENC);
                }};

                @Override
                public String getNamespaceURI(String prefix) {
                    return namespaceMap.get(prefix);
                }

                @Override
                public String getPrefix(String namespaceURI) {
                    return prefixMap.get(namespaceURI);
                }

                @Override
                public Iterator getPrefixes(String namespaceURI) {
                    return null;
                }
            });
            try {
                SimpleVariableResolver variableResolver = new SimpleVariableResolver();
                variableResolver.addVariable(new QName("Id"), refURI.substring(1));
                xpath.setXPathVariableResolver(variableResolver);
                Element result = (Element) xpath.evaluate("//xenc:EncryptedKey[@Id=$Id]", element.getOwnerDocument(), XPathConstants.NODE);
                XMLCipher cipher = XMLCipher.getInstance();
                cipher.init(XMLCipher.UNWRAP_MODE, this.decryptionKey);
                EncryptedKey ek = cipher.loadEncryptedKey(result);
                SecretKey secretKey = (SecretKey) cipher.decryptKey(ek, this.algorithm);
                log.debug("Got secret key: " + secretKey);
                return secretKey;
            } catch (XPathExpressionException e) {
                log.error("Error retrieving secret key: ", e);
            } catch (XMLEncryptionException e) {
                e.printStackTrace();  //To change body of catch statement use File | Settings | File Templates.
            }
        }
        return null;
    }

    private class SimpleVariableResolver implements XPathVariableResolver {
        private final Map vars = new HashMap<>();
        public void addVariable(QName name, Object value) {
            vars.put(name, value);
        }
        public Object resolveVariable(QName variableName) {
            return vars.get(variableName);
        }
    }
}