/*
 * Copyright 2017 Graz University of Technology EAAF-Core Components has been developed in a
 * cooperation between EGIZ, A-SIT Plus, A-SIT, and Graz University of Technology.
 *
 * Licensed under the EUPL, Version 1.2 or - as soon they will be approved by the European
 * Commission - subsequent versions of the EUPL (the "Licence"); You may not use this work except in
 * compliance with the Licence. You may obtain a copy of the Licence at:
 * https://joinup.ec.europa.eu/news/understanding-eupl-v12
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the Licence
 * is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the Licence for the specific language governing permissions and limitations under
 * the Licence.
 *
 * This product combines work with different licenses. See the "NOTICE" text file for details on the
 * various modules and licenses. The "NOTICE" text file is part of the distribution. Any derivative
 * works that you distribute must include a readable copy of the "NOTICE" text file.
*/

package at.gv.egiz.eaaf.modules.pvp2.idp.impl.builder;

import java.time.Instant;
import java.util.ArrayList;
import java.util.List;

import org.opensaml.core.criterion.EntityIdCriterion;
import org.opensaml.core.xml.config.XMLObjectProviderRegistrySupport;
import org.opensaml.core.xml.io.MarshallingException;
import org.opensaml.core.xml.io.UnmarshallingException;
import org.opensaml.core.xml.util.XMLObjectSupport;
import org.opensaml.saml.common.xml.SAMLConstants;
import org.opensaml.saml.criterion.EntityRoleCriterion;
import org.opensaml.saml.criterion.ProtocolCriterion;
import org.opensaml.saml.metadata.resolver.impl.PredicateRoleDescriptorResolver;
import org.opensaml.saml.saml2.core.Assertion;
import org.opensaml.saml.saml2.core.EncryptedAssertion;
import org.opensaml.saml.saml2.core.Issuer;
import org.opensaml.saml.saml2.core.NameIDType;
import org.opensaml.saml.saml2.core.RequestAbstractType;
import org.opensaml.saml.saml2.core.Response;
import org.opensaml.saml.saml2.encryption.Encrypter;
import org.opensaml.saml.saml2.encryption.Encrypter.KeyPlacement;
import org.opensaml.saml.saml2.metadata.SPSSODescriptor;
import org.opensaml.saml.security.impl.MetadataCredentialResolver;
import org.opensaml.security.credential.UsageType;
import org.opensaml.security.criteria.UsageCriterion;
import org.opensaml.security.x509.X509Credential;
import org.opensaml.xmlsec.SecurityConfigurationSupport;
import org.opensaml.xmlsec.encryption.support.DataEncryptionParameters;
import org.opensaml.xmlsec.encryption.support.EncryptionException;
import org.opensaml.xmlsec.encryption.support.KeyEncryptionParameters;
import org.opensaml.xmlsec.keyinfo.KeyInfoCredentialResolver;
import org.opensaml.xmlsec.keyinfo.KeyInfoGeneratorFactory;
import org.opensaml.xmlsec.keyinfo.impl.BasicProviderKeyInfoCredentialResolver;
import org.opensaml.xmlsec.keyinfo.impl.KeyInfoProvider;
import org.opensaml.xmlsec.keyinfo.impl.provider.DSAKeyValueProvider;
import org.opensaml.xmlsec.keyinfo.impl.provider.InlineX509DataProvider;
import org.opensaml.xmlsec.keyinfo.impl.provider.RSAKeyValueProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.w3c.dom.Element;

import at.gv.egiz.eaaf.core.api.idp.IConfiguration;
import at.gv.egiz.eaaf.modules.pvp2.PvpConstants;
import at.gv.egiz.eaaf.modules.pvp2.api.metadata.IPvp2MetadataProvider;
import at.gv.egiz.eaaf.modules.pvp2.exception.SamlSigningException;
import at.gv.egiz.eaaf.modules.pvp2.idp.exception.InvalidAssertionEncryptionException;
import at.gv.egiz.eaaf.modules.pvp2.impl.utils.Saml2Utils;
import net.shibboleth.shared.component.ComponentInitializationException;
import net.shibboleth.shared.resolver.CriteriaSet;
import net.shibboleth.shared.resolver.ResolverException;

/**
 * Authentication response builder.
 *
 * @author tlenz
 *
 */
public class AuthResponseBuilder {

  private static final Logger log = LoggerFactory.getLogger(AuthResponseBuilder.class);

  /**
   * Build PVP2 S-Profile authentication response.
   *
   * @param metadataProvider Service-Provider metadata
   * @param issuerEntityID   IDP entityId
   * @param req              current pending request
   * @param date             Timestamp
   * @param assertion        PVP2 S-Profil Assertion
   * @param authConfig       {@link IConfiguration}
   * @return PVP2 S-Profile authentication response
   * @throws InvalidAssertionEncryptionException In case of an error
   */
  public static Response buildResponse(final IPvp2MetadataProvider metadataProvider,
      final String issuerEntityID, final RequestAbstractType req, final Instant date,
      final Assertion assertion, IConfiguration authConfig)
      throws InvalidAssertionEncryptionException {
    final Response authResponse = Saml2Utils.createSamlObject(Response.class);

    final Issuer nissuer = Saml2Utils.createSamlObject(Issuer.class);

    nissuer.setValue(issuerEntityID);
    nissuer.setFormat(NameIDType.ENTITY);
    authResponse.setIssuer(nissuer);
    authResponse.setInResponseTo(req.getID());

    // set responseID
    final String remoteSessionID = Saml2Utils.getSecureIdentifier();
    authResponse.setID(remoteSessionID);

    // SAML2 response required IssueInstant
    authResponse.setIssueInstant(date);

    authResponse.setStatus(Saml2Utils.getSuccessStatus());

    // check, if metadata includes an encryption key
    final X509Credential encryptionCredentials = resolveEncryptionCredential(req, metadataProvider);

    if (encryptionCredentials != null
        && authConfig.getBasicConfigurationBoolean(
            PvpConstants.CONFIG_PROPERTY_PVP2_ENABLE_ENCRYPTION, true)) {
      authResponse.getEncryptedAssertions().add(
          doEncryption(assertion, encryptionCredentials, authConfig));

    } else {
      authResponse.getAssertions().add(assertion);

    }

    return authResponse;
  }

  private static EncryptedAssertion doEncryption(Assertion assertion,
      X509Credential encryptionCredentials, IConfiguration authConfig)
      throws InvalidAssertionEncryptionException {
    try {
      final String keyEncAlg = Saml2Utils.getKeyOperationAlgorithmFromCredential(
          encryptionCredentials,
          authConfig.getBasicConfiguration(
              PvpConstants.CONFIG_PROP_SEC_ENCRYPTION_KEY_RSA_ALG,
              PvpConstants.DEFAULT_ASYM_ENCRYPTION_METHODE_RSA),
          authConfig.getBasicConfiguration(
              PvpConstants.CONFIG_PROP_SEC_ENCRYPTION_KEY_EC_ALG,
              PvpConstants.DEFAULT_ASYM_ENCRYPTION_METHODE_EC));

      final DataEncryptionParameters dataEncParams = new DataEncryptionParameters();
      dataEncParams.setAlgorithm(authConfig.getBasicConfiguration(
          PvpConstants.CONFIG_PROP_SEC_ENCRYPTION_DATA, PvpConstants.DEFAULT_SYM_ENCRYPTION_METHODE));

      final List<KeyEncryptionParameters> keyEncParamList = new ArrayList<>();
      final KeyEncryptionParameters keyEncParam = new KeyEncryptionParameters();
      keyEncParam.setEncryptionCredential(encryptionCredentials);
      keyEncParam.setAlgorithm(keyEncAlg);

      final KeyInfoGeneratorFactory kigf =
          SecurityConfigurationSupport.getGlobalEncryptionConfiguration()
              .getKeyTransportKeyInfoGeneratorManager().getDefaultManager().getFactory(encryptionCredentials);
      keyEncParam.setKeyInfoGenerator(kigf.newInstance());
      keyEncParamList.add(keyEncParam);

      final Encrypter samlEncrypter = new Encrypter(dataEncParams, keyEncParamList);
      samlEncrypter.setKeyPlacement(KeyPlacement.PEER);

      final Element assertionElement = XMLObjectProviderRegistrySupport.getMarshallerFactory()
          .getMarshaller(assertion).marshall(assertion);
      assertionElement.setAttributeNS("http://www.w3.org/2000/xmlns/", "xmlns:xs", "http://www.w3.org/2001/XMLSchema");

      return samlEncrypter.encrypt((Assertion)
          XMLObjectSupport.getUnmarshaller(assertionElement).unmarshall(assertionElement));

    } catch (final EncryptionException | SamlSigningException | MarshallingException | UnmarshallingException e1) {
      log.warn("Can not encrypt the PVP2 assertion", e1);
      throw new InvalidAssertionEncryptionException();

    }

  }

  private static X509Credential resolveEncryptionCredential(RequestAbstractType req,
      IPvp2MetadataProvider metadataProvider) throws InvalidAssertionEncryptionException {
    try {
      final List<KeyInfoProvider> keyInfoProvider = new ArrayList<>();
      keyInfoProvider.add(new DSAKeyValueProvider());
      keyInfoProvider.add(new RSAKeyValueProvider());
      keyInfoProvider.add(new InlineX509DataProvider());
      final KeyInfoCredentialResolver keyInfoCredentialResolver = new BasicProviderKeyInfoCredentialResolver(
          keyInfoProvider);

      final PredicateRoleDescriptorResolver roleDescriptorResolver = new PredicateRoleDescriptorResolver(
          metadataProvider);
      roleDescriptorResolver.setRequireValidMetadata(true);
      roleDescriptorResolver.initialize();

      final MetadataCredentialResolver mdCredResolver = new MetadataCredentialResolver();
      mdCredResolver.setRoleDescriptorResolver(roleDescriptorResolver);
      mdCredResolver.setKeyInfoCredentialResolver(keyInfoCredentialResolver);
      mdCredResolver.initialize();

      final CriteriaSet criteriaSet = new CriteriaSet();
      criteriaSet.add(new EntityIdCriterion(req.getIssuer().getValue()));
      criteriaSet.add(new ProtocolCriterion(SAMLConstants.SAML20P_NS));
      criteriaSet.add(new EntityRoleCriterion(SPSSODescriptor.DEFAULT_ELEMENT_NAME));
      criteriaSet.add(new UsageCriterion(UsageType.ENCRYPTION));

      return (X509Credential) mdCredResolver.resolveSingle(criteriaSet);

    } catch (final SecurityException | ComponentInitializationException | ResolverException e2) {
      log.warn("Can not extract the Assertion Encryption-Key from metadata", e2);
      throw new InvalidAssertionEncryptionException();

    }
  }
}