EncryptPartialMessageProcessor.java

/*
 * AdroitLogic UltraESB Enterprise Service Bus
 *
 * Copyright (c) 2010-2012 AdroitLogic Private Ltd. (http://adroitlogic.org). All Rights Reserved.
 *
 * GNU Affero General Public License Usage
 *
 * This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General
 * Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option)
 * any later version.
 *
 * This program 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 Affero General Public License for
 * more details.
 *
 * You should have received a copy of the GNU Affero General Public License along with this program (See LICENSE-AGPL.TXT).
 * If not, see http://www.gnu.org/licenses/agpl-3.0.html
 *
 * Commercial Usage
 *
 * Licensees holding valid UltraESB Commercial licenses may use this file in accordance with the UltraESB Commercial
 * License Agreement provided with the Software or, alternatively, in accordance with the terms contained in a written
 * agreement between you and AdroitLogic.
 *
 * If you are unsure which license is appropriate for your use, or have questions regarding the use of this file,
 * please contact AdroitLogic at info@adroitlogic.com
 */
/*
 * Modificato da Link.it (https://link.it) per supportare le seguenti funzionalità:
 * - firma e cifratura degli attachments
 * - cifratura con chiave simmetrica
 * - supporto CRL 
 * 
 * Copyright (c) 2011-2024 Link.it srl (https://link.it). 
 * 
 */

package org.openspcoop2.security.message.soapbox;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.Security;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import javax.activation.DataHandler;
import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.xml.namespace.QName;
import javax.xml.soap.AttachmentPart;
import javax.xml.soap.MimeHeader;
import javax.xml.soap.MimeHeaders;
import javax.xml.soap.SOAPElement;
import javax.xml.soap.SOAPException;
import javax.xml.soap.SOAPHeaderElement;

import org.adroitlogic.soapbox.CryptoSupport;
import org.adroitlogic.soapbox.CryptoUtil;
import org.adroitlogic.soapbox.EncryptionRequest;
import org.adroitlogic.soapbox.InvalidMessageDataException;
import org.adroitlogic.soapbox.InvalidOptionException;
import org.adroitlogic.soapbox.MessageSecurityContext;
import org.adroitlogic.soapbox.Processor;
import org.adroitlogic.soapbox.SBConstants;
import org.adroitlogic.soapbox.SecurityFailureException;
import org.apache.xml.security.encryption.EncryptedData;
import org.apache.xml.security.encryption.XMLCipher;
import org.apache.xml.security.encryption.XMLEncryptionException;
import org.apache.xml.security.keys.KeyInfo;
import org.openspcoop2.message.OpenSPCoop2SoapMessage;
import org.openspcoop2.security.message.constants.WSSAttachmentsConstants;
import org.openspcoop2.utils.LoggerWrapperFactory;
import org.openspcoop2.utils.io.Base64Utilities;
import org.slf4j.Logger;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import com.sun.xml.wss.core.EncryptedDataHeaderBlock;
import com.sun.xml.wss.swa.MimeConstants;

/**
 * EncryptPartialMessageProcessor
 *
 * Author of the original AdroitLogic code:
 * @author asankha
 *
 * Authors of the Link.it modification to the code:
 * @author Andrea Poli (apoli@link.it)
 * @author $Author$
 * @version $Rev$, $Date$
 */
public class EncryptPartialMessageProcessor implements Processor {

	private static final Logger logger = LoggerWrapperFactory.getLogger(EncryptPartialMessageProcessor.class);


	private List<QName> elements;
	private List<Boolean> elementsEncryptContent;
	private Map<AttachmentPart, Boolean> attachments;
	private OpenSPCoop2SoapMessage message;
	public void setMessage(OpenSPCoop2SoapMessage message) {
		this.message = message;
	}
	private String actor;
	private boolean mustUnderstand;
	public void setActor(String actor) {
		this.actor = actor;
	}
	public void setMustUnderstand(boolean mustUnderstand) {
		this.mustUnderstand = mustUnderstand;
	}

	public EncryptPartialMessageProcessor() {
		this.elements = new ArrayList<QName>();
		this.elementsEncryptContent = new ArrayList<Boolean>();
		this.attachments = new HashMap<AttachmentPart, Boolean>();
	}

	public void addElementToEncrypt(QName element , boolean content) {
		this.elements.add(element);
		this.elementsEncryptContent.add(content);
	}

	public void addAttachmentToEncrypt(AttachmentPart part, boolean contentOnly) {
		this.attachments.put(part, contentOnly);
	}


	@Override
	public void process(org.adroitlogic.soapbox.SecurityConfig secConfig, MessageSecurityContext msgSecCtx) {

		Document doc = msgSecCtx.getDocument();

		// ensure existence of the wsse:Security header, and create one if none exists
        Element wsseSecurityElem = null;
        try{
        	wsseSecurityElem = WSSUtils.getWSSecurityHeader(msgSecCtx.getDocument(), this.actor, this.mustUnderstand);
        }catch(Exception e){
			throw new SecurityFailureException(e.getMessage(), e);
		}
		//Element wsseSecurityElem = CryptoUtil.getWSSecurityHeader(doc);
		
		// we will not encrypt an already encrypted document
		if (CryptoUtil.getFirstChildOrNull(wsseSecurityElem, SBConstants.XENC, "EncryptedKey") != null) {
			throw new InvalidMessageDataException("Message is already encrypted");
		}

		Element env = doc.getDocumentElement();

		String secTokenRef = CryptoUtil.getRandomId();

		// by default Body is encrypted
		if(this.elements.isEmpty() && this.attachments.isEmpty()){
			this.elements.add(new QName(env.getNamespaceURI(), "Body"));
			this.elementsEncryptContent.add(true);
		}


		// Process element e attachments
		try {

			processElements(msgSecCtx, secTokenRef);
			processAttachments(msgSecCtx);

		} catch (Exception e) {
			throw new SecurityFailureException("Error encrypting an element or an attachment", e);
		}
		
		
		
		// Process KeyInstance
		
		Cipher cipher = null;
		String cipherValue = null;

		EncryptionRequest encReq = msgSecCtx.getEncryptionRequest();

		try {
			SoapBoxSecurityConfig secConfigOpenSPCoop = (SoapBoxSecurityConfig)secConfig;

			Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
			cipher = CryptoSupport.getInstance().getCipherInstance(encReq.getEncryptionAlgoURI());

			byte[] encKey = null;
			SecretKey encKeyObject = null;
			try {
				encKeyObject = encReq.getEphemeralKey();
				encKey = encReq.getEphemeralKey().getEncoded();
			} catch (NoSuchAlgorithmException ignore) { /*will/should not happen*/ }

			if( secConfigOpenSPCoop.isSymmetricSharedKey() ){
				
	            int blockSize = cipher.getBlockSize();
	            if(blockSize==0){
	            	blockSize = 8;
	            }
	            //System.out.println("cipher (Algoritmo["+cipher.getAlgorithm()+"]) blksize: " + blockSize);
					            
	            SecureRandom randomSecureRandom = SecureRandom.getInstance("SHA1PRNG");
	            byte[] iv = new byte[blockSize];
	            randomSecureRandom.nextBytes(iv);
	            IvParameterSpec ivParams = new IvParameterSpec(iv);
	            
	            //cipher.init(Cipher.WRAP_MODE, secConfigOpenSPCoop.getSymmetricKey(encReq.getCertAlias()), new IvParameterSpec(new byte[blockSize]));
	            cipher.init(Cipher.WRAP_MODE, secConfigOpenSPCoop.getSymmetricKey(encReq.getCertAlias()), ivParams);
	            
				cipherValue = Base64Utilities.encodeAsString(cipher.wrap(encKeyObject));
			}
			else{

				cipher.init(Cipher.ENCRYPT_MODE, secConfig.getTrustedCertificatesByAlias(encReq.getCertAlias())[0]);
				
				if(encKey==null) {
					throw new SecurityFailureException("EncKey is null");
				}
				
				if (cipher.getBlockSize() > 0 && cipher.getBlockSize() < encKey.length) {
					throw new SecurityFailureException("Public key algorithm too weak to encrypt symmetric key " +
							" - cipher block size : " + cipher.getBlockSize() + " encrypted bytes length : " + encKey.length);
				}

				cipherValue = Base64Utilities.encodeAsString(cipher.doFinal(encKey));
			}

		} catch (InvalidKeyException e) {
			e.printStackTrace(System.err);
			throw new SecurityFailureException("Error preparing cipher for encryption : " + encReq.getEncryptionAlgoURI());
		} catch (Exception e) {
			throw new SecurityFailureException("Failed to encrypt session key", e);
		} finally {
			CryptoSupport.getInstance().returnCipherInstance(cipher);
		}

		Element encryptedKeyElem = createEncryptedKey(doc, encReq, cipherValue, secConfig, msgSecCtx, secTokenRef);
		Element firstChild = CryptoUtil.getFirstElementChild(wsseSecurityElem);
		if (firstChild != null) {
			wsseSecurityElem.insertBefore(encryptedKeyElem, firstChild);
		} else {
			wsseSecurityElem.appendChild(encryptedKeyElem);
		}
		
	}

	private Element createEncryptedKey(Document doc, EncryptionRequest encReq,
			String cipherValue, org.adroitlogic.soapbox.SecurityConfig secConfig, MessageSecurityContext msgSecCtx, String referenceId) {

		// create EncryptedKey element, and append EncryptionMethod with algorithm used
		Element encryptedKeyElem = doc.createElementNS(SBConstants.XENC, "xenc:EncryptedKey");
		encryptedKeyElem.setAttribute("Id", referenceId);
		Element encryptionMethodElem = doc.createElementNS(SBConstants.XENC, "xenc:EncryptionMethod");
		encryptionMethodElem.setAttribute("Algorithm", encReq.getEncryptionAlgoURI());
		encryptedKeyElem.appendChild(encryptionMethodElem);

		// create and attach the keyinfo element
		SoapBoxSecurityConfig securityConfigOpenSPCoop = (SoapBoxSecurityConfig) secConfig;
		if(securityConfigOpenSPCoop.isSymmetricSharedKey()){
			encryptedKeyElem.appendChild(SymmetricCryptoUtils.createKeyInfoElement(doc, encReq, msgSecCtx, secConfig));
		}else{
			//encryptedKeyElem.appendChild(CryptoUtil.createKeyInfoElement(doc, encReq, msgSecCtx, secConfig));
			encryptedKeyElem.appendChild(WSSUtils.createKeyInfoElement(doc, encReq, msgSecCtx, secConfig));
		}

		// create CipherData element and store the encrypted cipher value
		Element cipherDataElem = doc.createElementNS(SBConstants.XENC, "xenc:CipherData");
		Element cipherValueElem = doc.createElementNS(SBConstants.XENC, "xenc:CipherValue");
		cipherValueElem.setTextContent(cipherValue);
		cipherDataElem.appendChild(cipherValueElem);
		encryptedKeyElem.appendChild(cipherDataElem);

		// crate ReferenceList element and store encrypted element IDs
		Element referenceListElem = doc.createElementNS(SBConstants.XENC, "xenc:ReferenceList");
		for (String id : msgSecCtx.getEncryptedReferenceList()) {
			Element dataReferenceElem = doc.createElementNS(SBConstants.XENC, "xenc:DataReference");
			dataReferenceElem.setAttribute("URI", "#" + id);
			referenceListElem.appendChild(dataReferenceElem);
		}
		encryptedKeyElem.appendChild(referenceListElem);

		return encryptedKeyElem;
	}

	private void processElements(MessageSecurityContext msgSecCtx,
			String referenceId) throws Exception {

		Document doc = msgSecCtx.getDocument(); 
		Element env = doc.getDocumentElement();
		// TODO
		int index = 0;
		for(QName name : this.elements) {
			Element encElement = CryptoUtil.getFirstChild(env, name.getNamespaceURI(), name.getLocalPart());
			
			// L'attributo wsu:Id non serve nella encryption
//			String encId = encElement.getAttributeNS(SBConstants.WSU, "Id");
//			if (encId == null || encId.length() == 0) {
//				encId = encElement.getAttribute("Id");
//			}
//			if (encId == null || encId.length() == 0) {
//				encId = CryptoUtil.getRandomId();
//				encElement.setAttributeNS(SBConstants.WSU, "wsu:Id", encId);
//				CryptoUtil.setWsuId(encElement, encId);
//			}

			EncryptionRequest encReq = msgSecCtx.getEncryptionRequest();
			String symEncAlgo = encReq.getSymmetricKeyAlgoURI();
			XMLCipher xmlCipher = null;
			try {
				//xmlCipher = XMLCipher.getInstance(symEncAlgo);       	
				xmlCipher = CryptoSupport.getInstance().getXMLCipher(symEncAlgo);

				xmlCipher.init(XMLCipher.ENCRYPT_MODE, msgSecCtx.getEncryptionRequest().getEphemeralKey());
				EncryptedData encData = xmlCipher.getEncryptedData();
				String encEltId = CryptoUtil.getRandomId();
				encData.setId(encEltId);

				KeyInfo keyInfo = new KeyInfo(doc);
				Element securityTokenReferenceElem = doc.createElementNS(SBConstants.WSSE, "wsse:SecurityTokenReference");
				securityTokenReferenceElem.setAttributeNS(SBConstants.XMLNS, "xmlns:wsse", SBConstants.WSSE);
				Element referenceElem = doc.createElementNS(SBConstants.WSSE, "wsse:Reference");
				referenceElem.setAttribute("URI", "#" + referenceId);
				securityTokenReferenceElem.appendChild(referenceElem);

				keyInfo.addUnknownElement(securityTokenReferenceElem);
				encData.setKeyInfo(keyInfo);

				xmlCipher.doFinal(encElement.getOwnerDocument(), encElement, this.elementsEncryptContent.get(index++));
				msgSecCtx.addEncryptedReference(encEltId);

			} catch (XMLEncryptionException e) {
				throw new InvalidOptionException("Unsupported algorithm : " + symEncAlgo, e);
			} finally {
				try{
					CryptoSupport.getInstance().returnXMLCipherInstance(symEncAlgo, xmlCipher);
				}catch(Exception e){
					EncryptPartialMessageProcessor.logger.error(e.getMessage(),e);
				}
			}
		}
	}

	private static byte[] serializeHeaders(List<MimeHeader> mhs) throws Exception {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();

		try {

			StringBuilder line = new StringBuilder();

			for (MimeHeader mh : mhs) {

				String name = mh.getName();
				String vlue = mh.getValue();

				line.append(name);
				line.append(":");
				line.append(vlue);
				line.append("\r\n");

			}

			line.append("\r\n");
			byte[] b = line.toString().getBytes("US-ASCII");
			baos.write(b, 0, b.length);

		} catch (Exception e) {
			throw new Exception(e);
		}

		return baos.toByteArray();
	}

	private void processAttachments(MessageSecurityContext msgSecCtx) throws Exception {

		if(this.attachments.size()<=0){
			return;
		}

		Cipher _attachmentEncryptor = null;
		try {
			EncryptionRequest encReq = msgSecCtx.getEncryptionRequest();

			String encAlgo = encReq.getSymmetricKeyAlgoURI();

			_attachmentEncryptor = CryptoSupport.getInstance().getCipherInstance(encAlgo);
			_attachmentEncryptor.init(Cipher.ENCRYPT_MODE, encReq.getEphemeralKey());

			for(AttachmentPart p : this.attachments.keySet()) {

				boolean contentOnly = this.attachments.get(p);

				// create n push an ED

				EncryptedDataHeaderBlock edhb = new EncryptedDataHeaderBlock();

				String id = CryptoUtil.getRandomId();

				edhb.setId(id);
				edhb.setType( (contentOnly ?  WSSAttachmentsConstants.ATTACHMENT_CONTENT_ONLY_URI : WSSAttachmentsConstants.ATTACHMENT_COMPLETE_URI));
				edhb.setMimeType(p.getContentType());
				String uri = p.getContentId();
				if (uri != null) {           
					if(uri.startsWith("<")){
						uri = "cid:" + uri.substring(1, uri.length()-1);
					}else{
						uri = "cid:" + uri;
					}
				} else {
					uri = p.getContentLocation();
				}

				edhb.getCipherReference(true, uri);
				edhb.setEncryptionMethod(encAlgo);
				edhb.addTransform((contentOnly ?  WSSAttachmentsConstants.ATTACHMENT_CIPHERTEXT_TRANSFORM_URI : WSSAttachmentsConstants.ATTACHMENT_COMPLETE_TRANSFORM_URI));

				//System.out.println(" --PRIMA ENCRYPT --");
				//System.out.println(org.openspcoop2.pdd.logger.Dump.dumpMessage(this.message, true));
    							
				AttachmentPart encPart = EncryptPartialMessageProcessor.encryptAttachment(p, contentOnly, _attachmentEncryptor, this.message.createAttachmentPart());

				//System.out.println(" --DOPO ENCRYPT --");
				//System.out.println(org.openspcoop2.pdd.logger.Dump.dumpMessage(this.message, true));
				
//				this.attachments.remove(p); // Concurrent Modification
				MimeHeaders mhs = new MimeHeaders();
				mhs.addHeader(MimeConstants.CONTENT_ID, p.getContentId());
				this.message.removeAttachments(mhs);
				this.message.addAttachmentPart(encPart);
				
				msgSecCtx.addEncryptedReference(edhb.getId());
				//Element wssHeader = CryptoUtil.getWSSecurityHeader(msgSecCtx.getDocument());
				SOAPHeaderElement wssHeader = WSSUtils.getWSSecurityHeader(this.message, this.actor, this.mustUnderstand);
				
				SOAPElement elementToInsert = edhb.getAsSoapElement();
				wssHeader.addChildElement(elementToInsert);
				
				/*
				SOAPElement soapWssHeader = (SOAPElement) wssHeader;
				SecurityHeader _secHeader = new SecurityHeader(soapWssHeader) ;
				_secHeader.appendChild(edhb);*/            
			}
			this.attachments.clear();
		}
		finally {
			CryptoSupport.getInstance().returnCipherInstance(_attachmentEncryptor);
		}
	}

	private static AttachmentPart encryptAttachment(AttachmentPart part, boolean contentOnly, Cipher cipher, AttachmentPart encPart) throws Exception {

		byte[] cipherInput = (contentOnly) ? EncryptPartialMessageProcessor.getBytesFromAttachments(part.getDataHandler()) : EncryptPartialMessageProcessor.getCipherInput(part); 
		byte[] cipherOutput = cipher.doFinal(cipherInput);
		
		byte[] iv = cipher.getIV();
		byte[] encryptedBytes = new byte[iv.length + cipherOutput.length];

		System.arraycopy(iv, 0, encryptedBytes, 0, iv.length);
		System.arraycopy(cipherOutput, 0, encryptedBytes, iv.length, cipherOutput.length);

		int cLength  = encryptedBytes.length;
		String cType = MimeConstants.APPLICATION_OCTET_STREAM_TYPE;
		String uri   = part.getContentId();
		//Step 9 and 10.SWA spec.

		if (uri != null){
			encPart.setMimeHeader(MimeConstants.CONTENT_ID, uri);
		}else {
			uri = part.getContentLocation();
			if (uri != null){
				encPart.setMimeHeader(MimeConstants.CONTENT_LOCATION, uri);
			}
		}
		encPart.setContentType(cType);
		encPart.setMimeHeader(MimeConstants.CONTENT_LENGTH, Integer.toString(cLength));
		encPart.setMimeHeader(MimeConstants.CONTENT_TRANSFER_ENCODING, "base64");

		EncryptedAttachmentDataHandler dh = new EncryptedAttachmentDataHandler(new EncryptedAttachmentDataSource(encryptedBytes));
		encPart.setDataHandler(dh);

		cipherInput = (contentOnly) ? EncryptPartialMessageProcessor.getBytesFromAttachments(encPart.getDataHandler()) : EncryptPartialMessageProcessor.getCipherInput(encPart); 

		return encPart;
	}

	private static byte[] getCipherInput(AttachmentPart part) throws Exception,
	SOAPException, IOException {
		byte[] cipherInput;
		byte[] headers = EncryptPartialMessageProcessor.getAttachmentHeaders(part.getAllMimeHeaders());
		byte[] content = EncryptPartialMessageProcessor.getBytesFromAttachments(part.getDataHandler());

		cipherInput = new byte[headers.length+content.length];

		System.arraycopy(headers, 0, cipherInput, 0, headers.length);
		System.arraycopy(content, 0, cipherInput, headers.length, content.length);
		return cipherInput;
	}

	private static byte[] getAttachmentHeaders(Iterator<MimeHeader> mhItr) throws Exception {

		List<MimeHeader> mhs = new ArrayList<MimeHeader>();
		while (mhItr.hasNext()) mhs.add(mhItr.next());
		return EncryptPartialMessageProcessor.serializeHeaders(mhs);

	}

	private static byte[] getBytesFromAttachments(DataHandler dh) throws SOAPException, IOException {

		ByteArrayOutputStream baos = new ByteArrayOutputStream();  
		dh.writeTo(baos);
		return Base64Utilities.encodeAsString(baos.toByteArray()).getBytes("US-ASCII");

	}

	private static class EncryptedAttachmentDataHandler extends javax.activation.DataHandler {

		EncryptedAttachmentDataHandler(javax.activation.DataSource ds) {
			super(ds);
		}

		@Override
		public void writeTo(OutputStream os) throws java.io.IOException {
			((ByteArrayOutputStream) getDataSource().getOutputStream()).writeTo(os);
		}
	}

	private static class EncryptedAttachmentDataSource implements javax.activation.DataSource {
		byte[] datasource;

		EncryptedAttachmentDataSource(byte[] ds) {
			this.datasource = ds;
		}

		@Override
		public String getContentType() {
			return MimeConstants.APPLICATION_OCTET_STREAM_TYPE;
		}

		@Override
		public InputStream getInputStream() throws java.io.IOException {
			return new ByteArrayInputStream(this.datasource);
		}

		@Override
		public String getName() {
			return "Encrypted Attachment DataSource";
		}

		@Override
		public OutputStream getOutputStream() throws java.io.IOException {
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			baos.write(this.datasource, 0, this.datasource.length);
			return baos;
		}
	}

}