/*
 * Copyright 2012 AMG.lab, a Bull Group Company
 * 
 * Licensed 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.xlcloud.console.saml2;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.StringWriter;
import java.net.URLEncoder;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.zip.Deflater;
import java.util.zip.DeflaterOutputStream;

import org.joda.time.DateTime;
import org.opensaml.common.SAMLObjectBuilder;
import org.opensaml.common.SAMLVersion;
import org.opensaml.saml2.core.AuthnRequest;
import org.opensaml.saml2.core.Issuer;
import org.opensaml.saml2.core.LogoutRequest;
import org.opensaml.saml2.core.NameID;
import org.opensaml.saml2.core.SessionIndex;
import org.opensaml.xml.Configuration;
import org.opensaml.xml.ConfigurationException;
import org.opensaml.xml.XMLObject;
import org.opensaml.xml.XMLObjectBuilderFactory;
import org.opensaml.xml.io.MarshallingException;
import org.opensaml.xml.parse.ParserPool;
import org.opensaml.xml.util.Base64;
import org.opensaml.xml.util.XMLHelper;
import org.opensaml.xml.util.XMLObjectHelper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.w3c.dom.Element;
import org.xlcloud.console.context.IdentityContext;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;

/**
 * This class contains all the logic for SAML authentication
 * 
 * @author Jakub Wachowski, AMG.net
 */
public class Saml2ServiceProvider {

    private static Saml2ServiceProvider instance;

    //singleton as we need to cache AuthnRequests
    synchronized public static Saml2ServiceProvider getInstance() {
        if (instance == null) {
            instance = new Saml2ServiceProvider();
        }

        return instance;
    }

    private Saml2Configuration config;
    
    
    private static final Logger LOG = LoggerFactory.getLogger(Saml2ServiceProvider.class);
    
    /**
     * Caches info during authentication process
     */
    private Cache<String, AuthnInfoHolder> authnRequestsCache = CacheBuilder.newBuilder().expireAfterWrite(10, TimeUnit.MINUTES).build();

    /**
     * Constructor - bootstraps the OpenSAML library
     */
    private Saml2ServiceProvider() {
        try {
            if(LOG.isDebugEnabled()) {
                LOG.debug("OpenSAML library bootstrap started");
            }
            
            org.opensaml.DefaultBootstrap.bootstrap();
            
            if(LOG.isDebugEnabled()) {
                LOG.debug("OpenSAML library bootstrap finished");
            }
        } catch (ConfigurationException e) {
            LOG.error("Unable to bootstrap the OpenSAML library", e);
            throw new RuntimeException(e);
        }
        
        config = new Saml2Configuration("sp.xml", "idp.xml");
        
        
    }


    /**
     * Creates an AuthnRequest and redirects the UA to IDP for authentication
     * @param initialRequestUrl the URL to which user will be redirected after successful authentication
     * @return redirection URL with SAML request
     */
    public String getLoginRedirectUrlWithSamlRequest(String initialRequestUrl) {
        return config.getIdpLoginUrl() + "?SAMLRequest=" + deflateRequest(newAuthnRequest(initialRequestUrl));
    }

    /**
     * Creates an LogoutRequest and redirects the UA to IDP for Single Logout
     * @param identityContext contains SAML specific data, like IDP's session id
     * @return redirection URL with SAML request
     */
    public String getLogoutUrlWithSamlRequest(IdentityContext identityContext) {
        return config.getIdpLogoutUrl() + "?SAMLRequest=" + deflateRequest(newLogoutRequest(identityContext));
    }
    
    /**
     * Returns authentication info and removes it from cache (i.e. auth info can be read ONCE only). Returns null if there is no info for specified AuthRequest id
     * 
     * @param authnRequestId
     */
    AuthnInfoHolder getAuthnInfo(String authnRequestId) {
        if(authnRequestId == null || authnRequestId.isEmpty()) {
            return null;
        }
        
        AuthnInfoHolder authInfo =  authnRequestsCache.getIfPresent(authnRequestId);
        if(authInfo != null) {
            authnRequestsCache.invalidate(authnRequestId);
        }
        
        return authInfo;
    }

    /**
     * @return
     */
    ParserPool getParserPool() {
        return config.getParserPool();
    }



    private AuthnRequest newAuthnRequest(String initialRequestUrl) {

        XMLObjectBuilderFactory builderFactory = Configuration.getBuilderFactory();
        @SuppressWarnings( "unchecked" )
        SAMLObjectBuilder<AuthnRequest> builder = (SAMLObjectBuilder<AuthnRequest>) builderFactory
                .getBuilder(AuthnRequest.DEFAULT_ELEMENT_NAME);
        DateTime issueInstant = new DateTime();

        AuthnRequest authRequest = builder.buildObject();
        authRequest.setForceAuthn(new Boolean(false));
        authRequest.setIsPassive(new Boolean(false));
        authRequest.setIssueInstant(issueInstant);
        authRequest.setProtocolBinding("urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST");
        authRequest.setAssertionConsumerServiceURL(config.getAssertionConsumerServiceURL());
        @SuppressWarnings( "unchecked" )
        Issuer issuer = ((SAMLObjectBuilder<Issuer>) builderFactory.getBuilder(Issuer.DEFAULT_ELEMENT_NAME)).buildObject();
        issuer.setValue(config.getServiceProviderName());
        authRequest.setIssuer(issuer);
        String authReqRandomId = UUID.randomUUID().toString();
        authRequest.setID(authReqRandomId);
        authRequest.setVersion(SAMLVersion.VERSION_20);
        authRequest.setDestination(config.getIdpLoginUrl());

        //need to cache the request to verify if response is correct
        authnRequestsCache.put(authReqRandomId, new AuthnInfoHolder(authRequest, initialRequestUrl));
        Saml2MessageLogger.log(authRequest);
        
        return authRequest;
    }

    /**
     * @param identityContext 
     * @return
     */
    private LogoutRequest newLogoutRequest(IdentityContext identityContext) {
        XMLObjectBuilderFactory builderFactory = Configuration.getBuilderFactory();
        @SuppressWarnings( "unchecked" )
        SAMLObjectBuilder<LogoutRequest> builder = (SAMLObjectBuilder<LogoutRequest>) builderFactory
                .getBuilder(LogoutRequest.DEFAULT_ELEMENT_NAME);
        
        LogoutRequest logoutRequest = builder.buildObject();
        logoutRequest.setIssueInstant(new DateTime());
        @SuppressWarnings( "unchecked" )
        Issuer issuer = ((SAMLObjectBuilder<Issuer>) builderFactory.getBuilder(Issuer.DEFAULT_ELEMENT_NAME)).buildObject();
        issuer.setValue(config.getServiceProviderName());
        logoutRequest.setIssuer(issuer);
        logoutRequest.setDestination(config.getIdpLogoutUrl());
        logoutRequest.setID(UUID.randomUUID().toString());
        logoutRequest.setVersion(SAMLVersion.VERSION_20);
        
        @SuppressWarnings( "unchecked" )
        NameID nameID = ((SAMLObjectBuilder<NameID>) builderFactory.getBuilder(NameID.DEFAULT_ELEMENT_NAME)).buildObject();
        nameID.setFormat(identityContext.getNameIdFormat());
        nameID.setValue(identityContext.getName());
        logoutRequest.setNameID(nameID);
        @SuppressWarnings( "unchecked" )
        SessionIndex sessionIndex = ((SAMLObjectBuilder<SessionIndex>) builderFactory.getBuilder(SessionIndex.DEFAULT_ELEMENT_NAME)).buildObject();
        sessionIndex.setSessionIndex(identityContext.getIdpSessionId());
        logoutRequest.getSessionIndexes().add(sessionIndex);
        
        Saml2MessageLogger.log(logoutRequest);
        return logoutRequest;
    }
    
    /**
     * @param request
     * @return
     */
    private String deflateRequest(XMLObject request) {

        Element authDOM;
        try {
            authDOM = XMLObjectHelper.marshall(request);
            Deflater deflater = new Deflater(Deflater.DEFLATED, true);
            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
            DeflaterOutputStream deflaterOutputStream = new DeflaterOutputStream(byteArrayOutputStream, deflater);
            
            StringWriter rspWrt = new StringWriter();
            XMLHelper.writeNode(authDOM, rspWrt);
            deflaterOutputStream.write(rspWrt.toString().getBytes());
            deflaterOutputStream.close();
            
            /* Encoding the compressed message */
            String encodedRequestMessage = Base64.encodeBytes(byteArrayOutputStream.toByteArray(), Base64.DONT_BREAK_LINES);
            return URLEncoder.encode(encodedRequestMessage, "UTF-8").trim();
        } catch (MarshallingException e) {
            LOG.error("Unable to encode the AuthnRequest", e);
            throw new RuntimeException(e);
        } catch (IOException e) {
            LOG.error("Unable to encode the AuthnRequest", e);
            throw new RuntimeException(e);
        }

    }
}
