ws-security-dom/src/main/java/org/apache/wss4j/dom/message/WSSecDKEncrypt.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.apache.wss4j.dom.message;

import java.util.ArrayList;
import java.util.List;

import javax.crypto.SecretKey;

import org.apache.wss4j.common.WSEncryptionPart;
import org.apache.wss4j.common.derivedKey.ConversationConstants;
import org.apache.wss4j.common.ext.WSSecurityException;
import org.apache.wss4j.common.token.Reference;
import org.apache.wss4j.common.token.SecurityTokenReference;
import org.apache.wss4j.common.util.KeyUtils;
import org.apache.wss4j.dom.WSConstants;
import org.apache.wss4j.dom.util.WSSecurityUtil;
import org.apache.xml.security.encryption.Serializer;
import org.apache.xml.security.keys.KeyInfo;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;

/**
 * Encrypts and signs parts of a message with derived keys derived from a
 * symmetric key. This symmetric key will be included as an EncryptedKey
 */
public class WSSecDKEncrypt extends WSSecDerivedKeyBase {

    private String symEncAlgo = WSConstants.AES_128;
    private int derivedKeyLength = -1;

    private List<Element> attachmentEncryptedDataElements;

    private Serializer encryptionSerializer;

    public WSSecDKEncrypt(WSSecHeader securityHeader) {
        super(securityHeader);
    }

    public WSSecDKEncrypt(Document doc) {
        super(doc);
    }

    @Override
    public void prepare(byte[] ephemeralKey) throws WSSecurityException {
        super.prepare(ephemeralKey);

        attachmentEncryptedDataElements = new ArrayList<>();
    }

    public Document build(byte[] ephemeralKey) throws WSSecurityException {

        //
        // Setup the encrypted key
        //
        prepare(ephemeralKey);
        //
        // prepend elements in the right order to the security header
        //
        prependDKElementToHeader();

        Element externRefList = encrypt();

        addAttachmentEncryptedDataElements();

        addExternalRefElement(externRefList);

        return getDocument();
    }

    public void addAttachmentEncryptedDataElements() {
        if (attachmentEncryptedDataElements != null) {
            for (int i = 0; i < attachmentEncryptedDataElements.size(); i++) {
                Element encryptedData = attachmentEncryptedDataElements.get(i);
                Element securityHeaderElement = getSecurityHeader().getSecurityHeaderElement();
                WSSecurityUtil.prependChildElement(securityHeaderElement, encryptedData);
            }
        }
    }

    public Element encrypt() throws WSSecurityException {
        if (getParts().isEmpty()) {
            getParts().add(WSSecurityUtil.getDefaultEncryptionPart(getDocument()));
        }

        return encryptForExternalRef(null, getParts());
    }

    /**
     * Encrypt one or more parts or elements of the message (external).
     *
     * This method takes a vector of <code>WSEncryptionPart</code> object that
     * contain information about the elements to encrypt. The method call the
     * encryption method, takes the reference information generated during
     * encryption and add this to the <code>xenc:Reference</code> element.
     * This method can be called after <code>prepare()</code> and can be
     * called multiple times to encrypt a number of parts or elements.
     *
     * The method generates a <code>xenc:Reference</code> element that <i>must</i>
     * be added to the SecurityHeader. See <code>addExternalRefElement()</code>.
     *
     * If the <code>dataRef</code> parameter is <code>null</code> the method
     * creates and initializes a new Reference element.
     *
     * @param dataRef A <code>xenc:Reference</code> element or <code>null</code>
     * @param references A list containing WSEncryptionPart objects
     * @return Returns the updated <code>xenc:Reference</code> element
     * @throws WSSecurityException
     */
    public Element encryptForExternalRef(Element dataRef, List<WSEncryptionPart> references)
        throws WSSecurityException {

        KeyInfo keyInfo = createKeyInfo();

        SecretKey key = getDerivedKey(symEncAlgo);

        Encryptor encryptor = new Encryptor();
        encryptor.setDoc(getDocument());
        encryptor.setSecurityHeader(getSecurityHeader());
        encryptor.setIdAllocator(getIdAllocator());
        encryptor.setCallbackLookup(callbackLookup);
        encryptor.setAttachmentCallbackHandler(attachmentCallbackHandler);
        encryptor.setStoreBytesInAttachment(storeBytesInAttachment);
        encryptor.setEncryptionSerializer(encryptionSerializer);
        encryptor.setWsDocInfo(getWsDocInfo());
        List<String> encDataRefs =
            encryptor.doEncryption(keyInfo, key, symEncAlgo, references, attachmentEncryptedDataElements);

        if (dataRef == null) {
            dataRef =
                getDocument().createElementNS(
                    WSConstants.ENC_NS, WSConstants.ENC_PREFIX + ":ReferenceList"
                );
        }
        return WSSecEncrypt.createDataRefList(getDocument(), dataRef, encDataRefs);
    }

    /**
     * Create a KeyInfo object.
     *
     * @return Returns the created KeyInfo object.
     */
    private KeyInfo createKeyInfo() {
        KeyInfo keyInfo = new KeyInfo(getDocument());
        SecurityTokenReference secToken = new SecurityTokenReference(getDocument());
        secToken.addWSSENamespace();
        if (addWSUNamespace) {
            secToken.addWSUNamespace();
        }
        Reference ref = new Reference(getDocument());
        ref.setURI("#" + getId());
        String ns =
            ConversationConstants.getWSCNs(getWscVersion())
                + ConversationConstants.TOKEN_TYPE_DERIVED_KEY_TOKEN;
        ref.setValueType(ns);
        secToken.setReference(ref);

        keyInfo.addUnknownElement(secToken.getElement());
        Element keyInfoElement = keyInfo.getElement();
        keyInfoElement.setAttributeNS(
            WSConstants.XMLNS_NS, "xmlns:" + WSConstants.SIG_PREFIX, WSConstants.SIG_NS
        );

        return keyInfo;
    }

    /**
     * Adds (prepends) the external Reference element to the Security header.
     *
     * The reference element <i>must</i> be created by the
     * <code>encryptForExternalRef() </code> method. The method adds the
     * reference element in the SecurityHeader.
     *
     * @param referenceList The external <code>enc:Reference</code> element
     */
    public void addExternalRefElement(Element referenceList) {
        if (referenceList != null) {
            Node node = getdktElement().getNextSibling();
            Element securityHeaderElement = getSecurityHeader().getSecurityHeaderElement();
            if (node != null && Node.ELEMENT_NODE == node.getNodeType()) {
                securityHeaderElement.insertBefore(referenceList, node);
            } else {
                // If (at this moment) DerivedKeyToken is the LAST element of
                // the security header
                securityHeaderElement.appendChild(referenceList);
            }
        }
    }


    /**
     * Set the symmetric encryption algorithm URI to use
     * @param algo the symmetric encryption algorithm URI to use
     */
    public void setSymmetricEncAlgorithm(String algo) {
        symEncAlgo = algo;
    }

    protected int getDerivedKeyLength() throws WSSecurityException {
        return derivedKeyLength > 0 ? derivedKeyLength : KeyUtils.getKeyLength(symEncAlgo);
    }

    public void setDerivedKeyLength(int keyLength) {
        derivedKeyLength = keyLength;
    }

    public List<Element> getAttachmentEncryptedDataElements() {
        return attachmentEncryptedDataElements;
    }

    public Serializer getEncryptionSerializer() {
        return encryptionSerializer;
    }

    public void setEncryptionSerializer(Serializer encryptionSerializer) {
        this.encryptionSerializer = encryptionSerializer;
    }
}