/*
 * [AwsHandlerResolver.java]
 * Modified from Amazon-supplied version
 */
package com.mindprod.aws.jax;

import com.mindprod.base64.Base64;
import com.mindprod.common18.EIO;
import org.w3c.dom.Element;
import org.w3c.dom.Node;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import javax.xml.namespace.QName;
import javax.xml.soap.SOAPBody;
import javax.xml.soap.SOAPException;
import javax.xml.soap.SOAPMessage;
import javax.xml.ws.handler.Handler;
import javax.xml.ws.handler.HandlerResolver;
import javax.xml.ws.handler.MessageContext;
import javax.xml.ws.handler.PortInfo;
import javax.xml.ws.handler.soap.SOAPHandler;
import javax.xml.ws.handler.soap.SOAPMessageContext;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.List;
import java.util.Set;
import java.util.TimeZone;

/**
 * AwsHandlerResolver does the timestamp, signing and Base64 encoding
 * Modified from Amazon-supplied version at
 * http://www.google.ca/url?sa=t&rct=j&q=awshandlerresolver&source=web&cd=1&ved=0CC4QFjAA&url=https%3A%2F%2Fforums
 * .aws.amazon.com%2Fservlet%2FJiveServlet%2Fdownload%2F40-33731-141120-2645%2Fawshandlerresolver
 * .java&ei=a3pRT5uVO4PSiALKrcm1Bg&usg=AFQjCNEfcxsl26mMTnIepsIQtY9xolEJHg&sig2=T7OgyLvQO4Lj9d6CfrHKHw
 * to use the smaller com.mindprod.base64.Base64 package.
 */
public class AwsHandlerResolver implements HandlerResolver
    {
    private final String awsSecretKey;

    public AwsHandlerResolver( String awsSecretKey )
        {
        this.awsSecretKey = awsSecretKey;
        }

    @SuppressWarnings( "unchecked" )
    public List<Handler> getHandlerChain( PortInfo portInfo )
        {
        List<Handler> handlerChain = new ArrayList<>();
        QName serviceQName = portInfo.getServiceName();
        if ( serviceQName.getLocalPart().equals( "AWSECommerceService" ) )
            {
            handlerChain.add( new AwsHandler( awsSecretKey ) );
            }
        return handlerChain;
        }

    private static class AwsHandler implements SOAPHandler<SOAPMessageContext>
        {
        private final byte[] secretBytes;

        public AwsHandler( String awsSecretKey )
            {
            secretBytes = stringToUtf8( awsSecretKey );
            }

        private void appendTextElement( Node node, String elementName, String elementText )
            {
            Element element = node.getOwnerDocument().createElement( elementName );
            element.setTextContent( elementText );
            node.appendChild( element );
            }

        private String getSignature( String operation, String timeStamp, byte[] secretBytes )
            {
            try
                {
                String toSign = operation + timeStamp;
                byte[] toSignBytes = stringToUtf8( toSign );
                Mac signer = Mac.getInstance( "HmacSHA256" );
                SecretKeySpec keySpec = new SecretKeySpec( secretBytes, "HmacSHA256" );
                signer.init( keySpec );
                signer.update( toSignBytes );
                byte[] signBytes = signer.doFinal();
                return new Base64().encode( signBytes );  // the signature
                }
            catch ( NoSuchAlgorithmException nsae )
                {
                throw new RuntimeException( "NoSuchAlgorithmException was thrown.", nsae );
                }
            catch ( InvalidKeyException ike )
                {
                throw new RuntimeException( "InvalidKeyException was thrown.", ike );
                }
            }

        private String getTimestamp()
            {
            Calendar calendar = Calendar.getInstance();
            SimpleDateFormat dateFormat = new SimpleDateFormat( "yyyy-MM-dd'T'HH:mm:ss'Z'" );
            dateFormat.setTimeZone( TimeZone.getTimeZone( "UTC" ) );
            return dateFormat.format( calendar.getTime() );
            }

        private byte[] stringToUtf8( String source )
            {
            return source != null ? source.getBytes( EIO.UTF8 ) : null;
            }

        public void close( MessageContext messagecontext )
            {
            }

        public Set<QName> getHeaders()
            {
            return null;
            }

        public boolean handleFault( SOAPMessageContext messagecontext )
            {
            return true;
            }

        public boolean handleMessage( SOAPMessageContext messagecontext )
            {
            Boolean outbound = ( Boolean ) messagecontext.get( MessageContext.MESSAGE_OUTBOUND_PROPERTY );
            if ( outbound )
                {
                try
                    {
                    SOAPMessage soapMessage = messagecontext.getMessage();
                    SOAPBody soapBody = soapMessage.getSOAPBody();
                    Node firstChild = soapBody.getFirstChild();
                    String timeStamp = getTimestamp();
                    String signature = getSignature( firstChild.getLocalName(), timeStamp, secretBytes );
                    appendTextElement( firstChild, "Signature", signature );
                    appendTextElement( firstChild, "Timestamp", timeStamp );
                    }
                catch ( SOAPException se )
                    {
                    throw new RuntimeException( "SOAPException was thrown.", se );
                    }
                }
            return true;
            }
        }
    }