package gov.va.med.cds.ws.saml.handler;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.util.List;
import java.util.zip.DataFormatException;

import javax.ws.rs.container.ContainerRequestContext;
import javax.ws.rs.core.Context;
import javax.ws.rs.core.HttpHeaders;

import org.apache.cxf.common.util.Base64Exception;
import org.apache.cxf.common.util.Base64Utility;
import org.apache.cxf.common.util.StringUtils;
import org.apache.cxf.jaxrs.utils.JAXRSUtils;
import org.apache.cxf.message.Message;
import org.apache.cxf.rs.security.saml.DeflateEncoderDecoder;
import org.apache.cxf.rs.security.saml.SamlHeaderInHandler;
import org.apache.wss4j.common.ext.WSSecurityException;
import org.apache.wss4j.dom.WSConstants;
import org.w3c.dom.Element;

import gov.va.med.cds.saml.SamlAssertionThreadLocal;
import gov.va.med.cds.saml.SamlAssertionXml;
import gov.va.med.cds.ws.saml.util.SamlUtility;

public class SamlAuthorizationHeaderInHandler extends SamlHeaderInHandler {
	
	private static final String SAML_AUTH = "SAML";
	
	@Context
	private HttpHeaders headers;
	
	@Override 
	public void filter(ContainerRequestContext context) {
			//unable to accept the optionally compressed encoded strings, can only accept non compressed base64 encoded assertions for the time being ...
			this.setUseDeflateEncoding(false);
			Message message = JAXRSUtils.getCurrentMessage();
		     
		    List<String> values = headers.getRequestHeader(HttpHeaders.AUTHORIZATION); 
		    if(values == null){
		    	 //authorization header is optional
		    	 return;
		    }
		    else if (values != null && values.size() == 1 && !values.get(0).startsWith(SAML_AUTH)) {
		         throwFault("If Authorization header is available it must use SAML profile", null);    
		    }
		        
		    String[] parts = StringUtils.split(values.get(0), " ");
		    if (parts.length != 2) {
		          throwFault("Authorization header is malformed", null);
		    }
		        
		    handleToken(message, parts[1]);	
	  }
	
	@Override
	protected void handleToken(Message message, String assertion) {
        // the assumption here is that saml:Assertion is directly available, however, it 
        // may be contained inside saml:Response or saml:ArtifactResponse/saml:Response
        if (assertion == null) {
            throwFault("SAML assertion is not available", null);
        }
        
        try {
            byte[] deflatedToken = Base64Utility.decode(assertion);
            InputStream is = useDeflateEncoding() 
                ? new DeflateEncoderDecoder().inflateToken(deflatedToken)
                : new ByteArrayInputStream(deflatedToken);
            Element token = readToken(message, is);
            if ("Assertion".equals(token.getLocalName()) ) {
            	String samlXML = SamlUtility.assertionToString(token);
            	SamlAssertionThreadLocal.set(new SamlAssertionXml(samlXML));
            }
		}
        catch (WSSecurityException e) {
        	 throwFault("Extracting SAML to XML string has failed", e);  
        }
        catch (Base64Exception ex) {
            throwFault("Base64 decoding has failed", ex);
        } 
        catch (DataFormatException exx) {
            throwFault("Encoded assertion can not be inflated", exx);
        }         
    }

}
