/////////////////////////////////////////////////////////////////////////
//
// © University of Southampton IT Innovation Centre, 2004
//
// Copyright in this library belongs to the IT Innovation Centre of
// 2 Venture Road, Chilworth Science Park, Southampton, SO16 7NP, UK.
//
// This software may not be used, sold, licensed, transferred, copied
// or reproduced in whole or in part in any manner or form or in or
// on any media by any person other than in accordance with the terms
// of the Licence Agreement supplied with the software, or otherwise
// without the prior written consent of the copyright owners.
//
// This software is distributed WITHOUT ANY WARRANTY, without even the
// implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE, except where stated in the Licence Agreement supplied with
// the software.
//
//      Created by:             Darren Marvin
//      Created date:           2003/09/20
//      Created for project:    GEMSS
//
/////////////////////////////////////////////////////////////////////////
//
//      Dependencies: None
//
/////////////////////////////////////////////////////////////////////////
//
//      Last commit info:       $Author: $
//                              $Date: $
//                              $Revision: $
//
/////////////////////////////////////////////////////////////////////////

package uk.ac.soton.itinnovation.gemss.transportmessaging.messaging.processing.providers.wssecurity;

import eu.gemss.components.security.token.*;
import eu.gemss.components.security.token.types.key.asymmetric.*;
import eu.gemss.components.security.token.types.pki.*;
import eu.gemss.components.transport.servicedescription.*;
import eu.gemss.components.transport.servicedescription.policytypes.*;
import java.io.*;
import java.security.PrivateKey;
import java.security.cert.*;
import java.util.*;
import java.util.logging.*;
import javax.xml.soap.*;
import org.apache.xml.security.c14n.*;
import org.w3c.dom.*;
import org.xml.sax.*;
import uk.ac.soton.itinnovation.gemss.security.context.token.types.*;
import uk.ac.soton.itinnovation.gemss.transportmessaging.messaging.*;
import uk.ac.soton.itinnovation.gemss.transportmessaging.messaging.SOAPMessage;
import uk.ac.soton.itinnovation.gemss.transportmessaging.messaging.context.*;
import uk.ac.soton.itinnovation.gemss.transportmessaging.messaging.processing.*;
import uk.ac.soton.itinnovation.gemss.transportmessaging.servicedescription.policytypes.*;
import uk.ac.soton.itinnovation.wssecit.*;
import uk.ac.soton.itinnovation.wssecit.policy.*;
import uk.ac.soton.itinnovation.wssecit.signatures.*;
import uk.ac.soton.itinnovation.wssecit.encoding.WSSecSOAPMessage;
import uk.ac.soton.itinnovation.wssecit.encoding.WSSecHeader;
import org.apache.xml.security.signature.XMLSignature;
import org.apache.xml.security.keys.KeyInfo;
import java.security.Principal;
import eu.gemss.components.transport.servicedescription.endpointtypes.*;



public class WSSecurityMessageProcessor implements MessageProcessor, WSSecX509TrustManager   {

  private static final String AUTHENTICATED_SUBJECT_DN = "eu.gemss.transport.authenticated.subject.dn";
  private static final String AUTHENTICATED_ISSUER_DN = "eu.gemss.transport.authenticated.issuer.dn";
    private Logger mLogger = Logger.getLogger("uk.ac.soton.itinnovation.gemss.transportmessaging.messaging.processing.providers.wssecurity.WSSecurityMessageProcessor");
    private SOAPMessage mMessage = null;
    private static final String DESCRIPTION = "WS-Security processor for enforcing security of messages";
    private byte[] mByteArray = null; //need to keep this for as long as the processor is needed.
    private ByteArrayInputStream mInByteStream = null;
    private BufferedInputStream mBufInStream = null;


    public String getDescription() {
        return DESCRIPTION;
    }



    public Message enforceMessage(ServiceDescription serviceDesc,Policy policyToImplement, Message message) throws MessagingException {
      Message retMsg = null;
      if(policyToImplement.getClass().getName().equals(new BodyX509SigVerifyWSPolicyImp().getClass().getName())) {
            retMsg = verify(serviceDesc,message);
        }
        else if(policyToImplement.getClass().getName().equals(new BodyX509SigCreateWSPolicyImp().getClass().getName())) {
            retMsg = create(serviceDesc,message);
        }
        else {
            mLogger.log(Level.SEVERE,getDescription() + " does not support the policy described as '" + policyToImplement.getPolicyTypeDescriptor().getDescription() + "'");
            mLogger.log(Level.SEVERE,"The policy class supplied is '" + policyToImplement.getClass().getName() + "'");
            throw new MessagingException(getDescription() + " does not support the policy described as '" + policyToImplement.getPolicyTypeDescriptor().getDescription() + "'");
        }

        return retMsg;
    }

    private Message verify(ServiceDescription serviceDesc,Message message) throws MessagingException {
        Document doc = null;
        InputStream inStream = null;
        try{

            mMessage = (SOAPMessage) message;

            //pull off the attachments
             Iterator attIterator = mMessage.getSOAPWireMessage().getAttachments();
             Map attachmentMap = new HashMap();
             List attachments = new ArrayList();
             while(attIterator.hasNext()) {
                 AttachmentPart aP = (AttachmentPart) attIterator.next();
                 attachmentMap.put(aP.getContentId(),aP.getDataHandler());
                 attachments.add(aP);
            }

            doc = mMessage.getDOMPart();
            //use the endpoint supplied in the service description
            ServiceEndpoint sP = serviceDesc.getServiceEndpoint();
            if(!(sP instanceof WSDLEndpoint)) {
              mLogger.log(Level.SEVERE,"Processor '" + getDescription() + "' requires WSDL services");
              throw new MessagingException("Processor '" + getDescription() + "' requires WSDL services");
            }
            WSDLEndpoint wsdlP = (WSDLEndpoint) sP;
            if(!wsdlP.endpointSet()) {
              mLogger.log(Level.SEVERE,"Processor '" + getDescription() + "' requires a valid endpoint as a SOAP actor value");
              throw new MessagingException("Processor '" + getDescription() + "' requires a valid endpoint as a SOAP actor value");
            }
            String actorURI = wsdlP.getEndpointURL().toExternalForm();

            //use the wssecit library
            BasicBodyX509PolicyEnforcer policyEnforcer = new BasicBodyX509PolicyEnforcer(actorURI,this);



            boolean success = policyEnforcer.process(doc,attachmentMap);

            //retrieve the authenticated DN, best efforts only
            Principal subjectDN = null;
            Principal issuerDN = null;
            try{
              WSSecSOAPMessage wsSecMsg = new WSSecSOAPMessage(doc);
              WSSecHeader[] header = wsSecMsg.getWSSecurityHeaders();
              for(int i=0;i<header.length && issuerDN==null && subjectDN==null;i++) {
                WSSecXMLSignature[] xmlSigs = header[i].getOrderedXMLSignatures();
                for(int j=0;j<xmlSigs.length && issuerDN==null && subjectDN==null;j++) {
                   XMLSignature xmlSig = new XMLSignature(xmlSigs[j].getSignatureElement(),"");
                   KeyInfo keyInfo = xmlSig.getKeyInfo();
                   X509Certificate cert = keyInfo.getX509Certificate();
                   subjectDN = cert.getSubjectDN();
                   issuerDN = cert.getIssuerDN();
                }
              }
              //set the value as a property on the SOAPMessage context
              if(subjectDN!=null) {
                mMessage.getMessageContext().setMessageProperty(AUTHENTICATED_SUBJECT_DN,subjectDN);
              }
              if(issuerDN!=null) {
                mMessage.getMessageContext().setMessageProperty(AUTHENTICATED_ISSUER_DN,issuerDN);
              }


            }
            catch(Exception ex) {
              mLogger.log(Level.WARNING,"Unable to obtain authenticated DN to set in session, ignoring",ex);
            }

            //create a new message
            attIterator = attachments.iterator();
            while(attIterator.hasNext()) {
                AttachmentPart aP = (AttachmentPart) attIterator.next();
                mMessage.getSOAPWireMessage().addAttachmentPart(aP);
            }
            SOAPMessage returnMsg = new SOAPMessage(mMessage.getSOAPWireMessage(),mMessage.getMessageContext());

            if(!success) {
                mLogger.log(Level.WARNING,"Message does not conform to the X509 Signature Body policy");
                throw new MessageProcessingException("Message does not conform to the X509 Signature Body policy");
            }


            return returnMsg;

        }
        catch(WSSecurityException ex) {
            mLogger.log(Level.SEVERE,"Unable to check supplied secure SOAP message",ex);
            //try to supply the server response message but make it clear that message isn't trusted.
            mLogger.log(Level.WARNING,"The message is not secure and is provided here for information only, it up to you if you trust it:");
            try{
                SOAPMessage soapMsg = (SOAPMessage) message;
                if(message!=null) {
                    //if it is a response from a server then it will contain a fault hopefully
                    javax.xml.soap.SOAPMessage wireMsg = soapMsg.getSOAPWireMessage();
                    SOAPBody soapBody = wireMsg.getSOAPPart().getEnvelope().getBody();
                    if(soapBody.hasFault()) {
                        SOAPFault soapFault = soapBody.getFault();
                        String faultCode = soapFault.getFaultCode();
                        String faultString = soapFault.getFaultString();
                        String faultDetail = soapFault.getDetail().toString();
                        mLogger.log(Level.WARNING,"SOAP fault code is '" + faultCode + "'");
                        mLogger.log(Level.WARNING,"SOAP fault string is '" + faultString + "'");
                        mLogger.log(Level.WARNING,"SOAP fault detail is '" + faultDetail + "'");
                        throw new ServiceException(faultString,faultDetail,faultCode);
                    }
                    else {
                        mLogger.log(Level.WARNING,"SOAP message is not secure",ex);
                        throw new MessageProcessingException(ex.getMessage());
                    }
                }
                else {
                    mLogger.log(Level.WARNING,"Unable to read SOAP message");
                    throw new MessageProcessingException("Unable to read SOAP message");
                }
            }
            catch(javax.xml.soap.SOAPException e) {
                try{
                    //just dump the input stream contents to the log
                    mLogger.log(Level.WARNING,"Could not create SOAP message from the input stream, input stream dumped out here:");
                    mLogger.log(Level.WARNING,new String("[" + org.apache.axis.utils.XMLUtils.DocumentToString(doc)+ "]"));
                    throw new MessageProcessingException("Could not create SOAP message from the input stream, input stream dumped out to log files");
                }
                catch(Exception e2) {
                    mLogger.log(Level.WARNING,"Could not obtain response stream, reason:",e2);
                }
            }
            catch(ClassCastException e) {
                //this is to catch the bug in the JAXM RI software
                try{
                    //just dump the input stream contents to the log
                    mLogger.log(Level.WARNING,"Could not create SOAP message from the input stream, input stream dumped out here:");
                    mLogger.log(Level.WARNING,new String("[" + org.apache.axis.utils.XMLUtils.DocumentToString(doc)+ "]"));
                    throw new MessageProcessingException("Could not create SOAP message from the input stream, input stream dumped out to log files");
                }
                catch(Exception e2) {
                    mLogger.log(Level.WARNING,"Could not obtain response stream, reason:",e2);
                }
            }
            catch(Exception e) {
                if(e instanceof MessageProcessingException)
                    throw (MessageProcessingException) e;
                mLogger.log(Level.WARNING,"Could not obtain response SOAP message, reason:",e);
                throw new MessageProcessingException("Unable to check supplied secure SOAP message, please consult the logs");
            }
            throw new MessageProcessingException("Unable to check supplied secure SOAP message, please consult the logs");
        }

        catch(java.io.IOException ex) {
            mLogger.log(Level.SEVERE,"Problem parsing supplied secure SOAP message",ex);
            throw new MessageProcessingException("Problem parsing supplied secure SOAP message, please consult the logs");
        }
        catch(javax.xml.parsers.ParserConfigurationException ex) {
            mLogger.log(Level.SEVERE,"Problem parsing supplied secure SOAP message",ex);
            throw new MessageProcessingException("Problem parsing supplied secure SOAP message, please consult the logs");
        }
        catch(javax.xml.soap.SOAPException ex) {
            mLogger.log(Level.SEVERE,"Problem parsing supplied secure SOAP message",ex);
            throw new MessageProcessingException("Problem parsing supplied secure SOAP message, please consult the logs");
        }
        catch(org.xml.sax.SAXException ex) {
            mLogger.log(Level.SEVERE,"Problem parsing supplied secure SOAP message, could be because web service is not available",ex);
            try{
                if(inStream.available()>0) {
                    byte[] response = new byte[inStream.available()];
                    mLogger.log(Level.SEVERE,"Part of the response is the following [" +  new String(response) + "]");
                }
            }
            catch(Exception e) {
                //fall through to next - not much we can do about it.
            }
            throw new MessageProcessingException("Problem parsing supplied secure SOAP message, please consult the logs");
        }

    }

    private Message create(ServiceDescription serviceDesc,Message message) throws MessagingException {
        try{

            //find out what the policy is exactly and do the necessary encoding.

            //first retrieve a private key and certificate from the security context
            SecurityContext secCtx = message.getMessageContext().getSecurityContext();
            PrivateKeyToken privateKeyToken = (PrivateKeyToken) secCtx.generateSecurityToken(
                    new AsymmetricKeyDescriptor(AsymmetricKeyDescriptor.PRIVATE_KEY,AsymmetricKeyDescriptor.CYPHER_ALGO_RSA));
            X509CertificateToken certToken = (X509CertificateToken) secCtx.generateSecurityToken(
                    new CertificateDescriptor(CertificateDescriptor.PKI_X509V3,AsymmetricKeyDescriptor.CYPHER_ALGO_RSA));
            PrivateKey privateKey = privateKeyToken.getKey();
            X509Certificate cert = certToken.getX509Certificate();
            //Create a DOM document instance
            //have to remove the attachments from the message temporarily to enable us to
            //create a DOM object

            SOAPMessage soapMsg = (SOAPMessage) message;


            //pull off the attachments
            Iterator attIterator = soapMsg.getSOAPWireMessage().getAttachments();
            Map attachmentMap = new HashMap();
            List attachments = new ArrayList();
            while(attIterator.hasNext()) {
                AttachmentPart aP = (AttachmentPart) attIterator.next();
                attachmentMap.put(aP.getContentId(),aP.getDataHandler());
                attachments.add(aP);
            }

            Document doc = soapMsg.getDOMPart();

            //use the endpoint supplied in the service description
            ServiceEndpoint sP = serviceDesc.getServiceEndpoint();
            if(!(sP instanceof WSDLEndpoint)) {
              mLogger.log(Level.SEVERE,"Processor '" + getDescription() + "' requires WSDL services");
              throw new MessagingException("Processor '" + getDescription() + "' requires WSDL services");
            }
            WSDLEndpoint wsdlP = (WSDLEndpoint) sP;
            if(!wsdlP.endpointSet()) {
              mLogger.log(Level.SEVERE,"Processor '" + getDescription() + "' requires a valid endpoint as a SOAP actor value");
              throw new MessagingException("Processor '" + getDescription() + "' requires a valid endpoint as a SOAP actor value");
            }
            String actorURI = wsdlP.getEndpointURL().toExternalForm();

            //use the wssecit library
            BasicBodyX509PolicyEnforcer policyEnforcer = new BasicBodyX509PolicyEnforcer(actorURI,privateKey,cert);
            Document securedDoc = policyEnforcer.enforce(doc,attachmentMap);

            //retrieve the authenticated DN, best efforts only
            Principal subjectDN = null;
            Principal issuerDN = null;
            try{
              WSSecSOAPMessage wsSecMsg = new WSSecSOAPMessage(securedDoc);
              WSSecHeader[] header = wsSecMsg.getWSSecurityHeaders();
              for(int i=0;i<header.length && issuerDN==null && subjectDN==null;i++) {
                WSSecXMLSignature[] xmlSigs = header[i].getOrderedXMLSignatures();
                for(int j=0;j<xmlSigs.length && issuerDN==null && subjectDN==null;j++) {
                   XMLSignature xmlSig = new XMLSignature(xmlSigs[j].getSignatureElement(),"");
                   KeyInfo keyInfo = xmlSig.getKeyInfo();
                   X509Certificate keyCert = keyInfo.getX509Certificate();
                   subjectDN = cert.getSubjectDN();
                   issuerDN = cert.getIssuerDN();
                }
              }
              //set the value as a property on the SOAPMessage context
              if(subjectDN!=null) {
                message.getMessageContext().setMessageProperty(AUTHENTICATED_SUBJECT_DN,subjectDN);
              }
              if(issuerDN!=null) {
                message.getMessageContext().setMessageProperty(AUTHENTICATED_ISSUER_DN,issuerDN);
              }

            }
            catch(Exception ex) {
              mLogger.log(Level.WARNING,"Unable to obtain authenticated DN to set in session, ignoring",ex);
            }

            Canonicalizer c14n = Canonicalizer.getInstance(Canonicalizer.ALGO_ID_C14N_WITH_COMMENTS);
            byte[] canonicalMessage = c14n.canonicalizeSubtree(securedDoc);
            InputSource is = new InputSource(new java.io.ByteArrayInputStream(canonicalMessage));
            SOAPMessage securedSoapMsg = new SOAPMessage(message.getMessageContext());
            javax.xml.transform.sax.SAXSource saxSource = new javax.xml.transform.sax.SAXSource(is);
            securedSoapMsg.getSOAPWireMessage().getSOAPPart().setContent(saxSource);

            attIterator = attachments.iterator();
            while(attIterator.hasNext()) {
                AttachmentPart aP = (AttachmentPart) attIterator.next();
                securedSoapMsg.getSOAPWireMessage().addAttachmentPart(aP);
            }

            return securedSoapMsg;

        }


        catch(javax.xml.soap.SOAPException ex) {
            mLogger.log(Level.SEVERE,"Unable to create a suitable SOAP message",ex);
            throw new MessageProcessingException("Unable to create a suitable SOAP message, please consult the logs");
        }

        catch(MessageContextException ex) {
            mLogger.log(Level.SEVERE,"Unable to establish a message context for the message",ex);
            throw new MessageProcessingException("Unable to establish a message context for the message, please consult the logs");
        }
        catch(SecurityTokenException ex) {
            mLogger.log(Level.SEVERE,"Unable to obtain a security token for the message",ex);
            throw new MessageProcessingException("Unable to obtain a security token for the message, please consult the logs");
        }

        catch(javax.xml.parsers.ParserConfigurationException ex) {
            mLogger.log(Level.SEVERE,"Unable to create a suitable SOAP message",ex);
            throw new MessageProcessingException("Unable to create a suitable SOAP message, please consult the logs");
        }

        catch(org.apache.xml.security.c14n.InvalidCanonicalizerException ex){
            mLogger.log(Level.SEVERE,"Unable to create a suitable SOAP message",ex);
            throw new MessageProcessingException("Unable to create a suitable SOAP message, please consult the logs");
        }
        catch(org.apache.xml.security.c14n.CanonicalizationException ex) {
            mLogger.log(Level.SEVERE,"Unable to create a suitable SOAP message",ex);
            throw new MessageProcessingException("Unable to create a suitable SOAP message, please consult the logs");
        }

        catch(java.io.IOException ex) {
            mLogger.log(Level.SEVERE,"Unable to create a suitable SOAP message",ex);
            throw new MessageProcessingException("Unable to create a suitable SOAP message, please consult the logs");
        }
        catch(org.xml.sax.SAXException ex) {
            mLogger.log(Level.SEVERE,"Unable to create a suitable SOAP message",ex);
            throw new MessageProcessingException("Unable to create a suitable SOAP message, please consult the logs");
        }

        catch(WSSecurityException ex) {
            mLogger.log(Level.SEVERE,"Unable to secure supplied SOAP message",ex);
            throw new MessageProcessingException("Unable to secure supplied SOAP message, please consult the logs");
        }

    }


    public boolean isTrusted(X509Certificate cert) {
        boolean trusted = false;
        try{
            if(cert==null) {
                mLogger.log(Level.SEVERE,"Cannot establish if supplied is trusted because supplied certificate is null, this is a bug - send a bug report");
            }
            else {
                X509Certificate[] chain = new X509Certificate[1];
                chain[0] = cert;
                trusted = mMessage.getMessageContext().getSecurityContext().isTrusted(new X509CertificateChainTokenImpl(chain,X509CertificateChainTokenImpl.AUTH_TYPE_RSA,false));
            }
            return trusted;
        }
        catch(MessageContextException ex) {
            mLogger.log(Level.SEVERE,"Could not establish trust for a certificate, please send a bug report",ex);
            return false;
        }
        catch(MessagingException ex) {
            mLogger.log(Level.SEVERE,"Could not establish trust for a certificate, please send a bug report",ex);
            return false;
        }

    }

}