package at.gv.egiz.eaaf.modules.pvp2.test;

import java.util.ArrayList;
import java.util.List;

import org.apache.xml.security.algorithms.JCEMapper;
import org.joda.time.DateTime;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;
import org.opensaml.core.xml.config.XMLObjectProviderRegistrySupport;
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.saml2.core.Assertion;
import org.opensaml.saml.saml2.core.AuthnRequest;
import org.opensaml.saml.saml2.core.EncryptedAssertion;
import org.opensaml.saml.saml2.core.Issuer;
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.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.KeyInfoGeneratorFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.Assert;

import at.gv.egiz.eaaf.core.api.idp.IConfiguration;
import at.gv.egiz.eaaf.core.exceptions.EaafException;
import at.gv.egiz.eaaf.core.exceptions.InvalidProtocolRequestException;
import at.gv.egiz.eaaf.core.impl.data.Pair;
import at.gv.egiz.eaaf.modules.pvp2.PvpConstants;
import at.gv.egiz.eaaf.modules.pvp2.api.credential.EaafX509Credential;
import at.gv.egiz.eaaf.modules.pvp2.api.metadata.IPvp2MetadataProvider;
import at.gv.egiz.eaaf.modules.pvp2.exception.CredentialsNotAvailableException;
import at.gv.egiz.eaaf.modules.pvp2.exception.Pvp2InternalErrorException;
import at.gv.egiz.eaaf.modules.pvp2.exception.Pvp2MetadataException;
import at.gv.egiz.eaaf.modules.pvp2.exception.SamlSigningException;
import at.gv.egiz.eaaf.modules.pvp2.impl.message.PvpSProfileRequest;
import at.gv.egiz.eaaf.modules.pvp2.impl.message.PvpSProfileResponse;
import at.gv.egiz.eaaf.modules.pvp2.impl.metadata.PvpMetadataResolverFactory;
import at.gv.egiz.eaaf.modules.pvp2.impl.opensaml.initialize.EaafOpenSaml3xInitializer;
import at.gv.egiz.eaaf.modules.pvp2.impl.utils.Saml2Utils;
import at.gv.egiz.eaaf.modules.pvp2.impl.validation.TrustEngineFactory;
import at.gv.egiz.eaaf.modules.pvp2.impl.verification.SamlVerificationEngine;
import at.gv.egiz.eaaf.modules.pvp2.test.dummy.DummyCredentialProvider;
import at.gv.egiz.eaaf.modules.pvp2.test.dummy.DummyMetadataProvider;
import net.shibboleth.utilities.java.support.xml.XMLParserException;


public abstract class AbstractSamlVerificationEngine {

  @Autowired
  private PvpMetadataResolverFactory metadataResolverFactory;
  @Autowired
  private SamlVerificationEngine verifyEngine;
  @Autowired
  protected DummyCredentialProvider credentialProvider;

  @Autowired DummyMetadataProvider metadataProvider;
  @Autowired IConfiguration authConfig;

  /**
   * JUnit class initializer.
   *
   * @throws Exception In case of an OpenSAML3 initialization error
   */
  @BeforeClass
  public static void classInitializer() throws Exception {
    EaafOpenSaml3xInitializer.eaafInitialize();

  }
  
  /**
   * Reset OpenSAML3.x JCEMapper to default.
   * 
   */
  @AfterClass
  public static void classCloser() {
    JCEMapper.setProviderId(null); 
   
  }
  
  protected abstract String getMetadataJunitJKeystore();

  protected abstract String getMetadataClassPathEntityPath();

  protected abstract String getAuthnRequestWithoutSigPath();

  protected abstract String getResponseWithSigPath();

  protected abstract String getResponseWithoutSigPath();


  @Test
  public void validateSamlRequestSuccess() throws SecurityException, Exception {

    final String authnReqPath = getAuthnRequestWithoutSigPath();
    final String metadataPath = getMetadataClassPathEntityPath();
    final String spEntityId = metadataPath;

    final Pair<AuthnRequest, IPvp2MetadataProvider> inputMsg =
        initializeAuthnRequest(spEntityId, metadataPath, authnReqPath,
            credentialProvider.getMetaDataSigningCredential());

    final PvpSProfileRequest msg = new PvpSProfileRequest(
        inputMsg.getFirst(),
        SAMLConstants.SAML2_POST_BINDING_URI);
    msg.setEntityID(spEntityId);

    verifyEngine.verify(msg,
        TrustEngineFactory.getSignatureKnownKeysTrustEngine(metadataProvider));

  }

  @Test
  public void validateSamlRequestWrongSignature() throws SecurityException, Exception {

    final String authnReqPath = getAuthnRequestWithoutSigPath();
    final String metadataPath = getMetadataJunitJKeystore();
    final String spEntityId = metadataPath;

    final Pair<AuthnRequest, IPvp2MetadataProvider> inputMsg =
        initializeAuthnRequest(spEntityId, metadataPath, authnReqPath,
            credentialProvider.getMetaDataSigningCredential());

    metadataProvider.addMetadataResolverIntoChain(inputMsg.getSecond());

    final PvpSProfileRequest msg = new PvpSProfileRequest(
        inputMsg.getFirst(),
        SAMLConstants.SAML2_POST_BINDING_URI);
    msg.setEntityID(spEntityId);

    try {
      verifyEngine.verify(msg,
          TrustEngineFactory.getSignatureKnownKeysTrustEngine(metadataProvider));
      org.junit.Assert.fail("Wrong signature not detected");

    } catch (final Exception e) {
      Assert.isInstanceOf(InvalidProtocolRequestException.class, e, "Wrong exceptionType");
      org.junit.Assert.assertEquals("Wrong errorcode", "internal.pvp.10", ((EaafException) e).getErrorId());

    }
  }

  @Test
  public void verifyResponseSuccessTest() throws Pvp2InternalErrorException, SecurityException, Exception {
    metadataProvider.runGarbageCollector();

    final String authnReqPath = getResponseWithoutSigPath();
    final String metadataPath = getMetadataClassPathEntityPath();
    final String spEntityId = metadataPath;

    final Pair<Response, IPvp2MetadataProvider> inputMsg =
        initializeResponse(spEntityId, metadataPath, authnReqPath,
            credentialProvider.getMetaDataSigningCredential());

    final PvpSProfileResponse msg = new PvpSProfileResponse(
        inputMsg.getFirst());
    msg.setEntityID(spEntityId);

    verifyEngine.verify(msg,
        TrustEngineFactory.getSignatureKnownKeysTrustEngine(metadataProvider));

  }

  @Test
  public void verifyResponseSuccessSecondTest()
      throws Pvp2InternalErrorException, SecurityException, Exception {

    final String authnReqPath = getResponseWithoutSigPath();
    final String metadataPath = getMetadataClassPathEntityPath();
    final String spEntityId = metadataPath;

    final Pair<Response, IPvp2MetadataProvider> inputMsg =
        initializeResponse(spEntityId, metadataPath, authnReqPath,
            credentialProvider.getMetaDataSigningCredential());

    verifyEngine.verifyIdpResponse(inputMsg.getFirst(),
        TrustEngineFactory.getSignatureKnownKeysTrustEngine(metadataProvider));

  }

  @Test
  public void verifySpResponse()
      throws Pvp2InternalErrorException, SecurityException, Exception {

    final String authnReqPath = getResponseWithoutSigPath();
    final String metadataPath = getMetadataClassPathEntityPath();
    final String spEntityId = metadataPath;

    final Pair<Response, IPvp2MetadataProvider> inputMsg =
        initializeResponse(spEntityId, metadataPath, authnReqPath,
            credentialProvider.getMetaDataSigningCredential());

    verifyEngine.verifySloResponse(inputMsg.getFirst(),
        TrustEngineFactory.getSignatureKnownKeysTrustEngine(metadataProvider));

  }

  @Test
  public void verifyResponseWithoutId() throws Pvp2InternalErrorException, SecurityException, Exception {

    final String authnReqPath = getResponseWithSigPath();
    final String metadataPath = getMetadataClassPathEntityPath();
    final String spEntityId = metadataPath;

    final Pair<Response, IPvp2MetadataProvider> inputMsg =
        initializeResponse(spEntityId, metadataPath, authnReqPath,
            credentialProvider.getMetaDataSigningCredential());

    final PvpSProfileResponse msg = new PvpSProfileResponse(
        inputMsg.getFirst());
    msg.setEntityID(spEntityId);

    try {
      verifyEngine.verify(msg,
          TrustEngineFactory.getSignatureKnownKeysTrustEngine(metadataProvider));
      org.junit.Assert.fail("Wrong XML schema not detected");

    } catch (final Exception e) {
      Assert.isInstanceOf(InvalidProtocolRequestException.class, e, "Wrong exceptionType");
      org.junit.Assert.assertEquals("Wrong errorcode", "internal.pvp.03", ((EaafException) e).getErrorId());

    }
  }

  @Test
  public void verifyResponseWrongTrust() throws Pvp2InternalErrorException, SecurityException, Exception {

    final String authnReqPath = getResponseWithoutSigPath();
    final String metadataPath = getMetadataJunitJKeystore();
    final String spEntityId = metadataPath;

    final Pair<Response, IPvp2MetadataProvider> inputMsg =
        initializeResponse(spEntityId, metadataPath, authnReqPath,
            credentialProvider.getMetaDataSigningCredential());

    final PvpSProfileResponse msg = new PvpSProfileResponse(
        inputMsg.getFirst());
    msg.setEntityID(spEntityId);

    try {
      verifyEngine.verify(msg,
          TrustEngineFactory.getSignatureKnownKeysTrustEngine(metadataProvider));
      org.junit.Assert.fail("No TrustedCert not detected");

    } catch (final Exception e) {
      Assert.isInstanceOf(InvalidProtocolRequestException.class, e, "Wrong exceptionType");
      org.junit.Assert.assertEquals("Wrong errorcode", "internal.pvp.10", ((EaafException) e).getErrorId());

    }
  }

  protected Pair<Response, IPvp2MetadataProvider> initializeResponse(String spEntityId, String metadataPath,
      String authnReqPath, EaafX509Credential credential)
          throws SamlSigningException, XMLParserException, UnmarshallingException, Pvp2MetadataException {
    final IPvp2MetadataProvider mdResolver = metadataResolverFactory.createMetadataProvider(
        metadataPath, null, "jUnit metadata resolver", null);

    final Response authnReq = (Response) XMLObjectSupport.unmarshallFromInputStream(
        XMLObjectProviderRegistrySupport.getParserPool(),
        AbstractSamlVerificationEngine.class.getResourceAsStream(authnReqPath));
    authnReq.setIssueInstant(DateTime.now());
    final Issuer issuer = Saml2Utils.createSamlObject(Issuer.class);
    issuer.setValue(spEntityId);
    authnReq.setIssuer(issuer);

    return Pair.newInstance(
        Saml2Utils.signSamlObject(authnReq, credential, true),
        mdResolver);
  }

  protected Pair<AuthnRequest, IPvp2MetadataProvider> initializeAuthnRequest(String spEntityId,
      String metadataPath, String authnReqPath, EaafX509Credential credential)
      throws SamlSigningException, CredentialsNotAvailableException,
      XMLParserException, UnmarshallingException, Pvp2InternalErrorException, Pvp2MetadataException {

    final IPvp2MetadataProvider mdResolver = metadataResolverFactory.createMetadataProvider(
        metadataPath, null, "jUnit metadata resolver", null);

    final AuthnRequest authnReq = (AuthnRequest) XMLObjectSupport.unmarshallFromInputStream(
        XMLObjectProviderRegistrySupport.getParserPool(),
        AbstractSamlVerificationEngine.class.getResourceAsStream(authnReqPath));
    authnReq.setIssueInstant(DateTime.now());
    final Issuer issuer = Saml2Utils.createSamlObject(Issuer.class);
    issuer.setValue(spEntityId);
    authnReq.setIssuer(issuer);

    return Pair.newInstance(
        Saml2Utils.signSamlObject(authnReq, credential, true),
        mdResolver);

  }

  protected static EncryptedAssertion doEncryption(Assertion assertion,
      X509Credential encryptionCredentials, IConfiguration authConfig)
      throws Exception {
    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);

      return samlEncrypter.encrypt(assertion);

    } catch (final EncryptionException | SamlSigningException e1) {
      throw new Exception(e1);

    }

  }


}