/*
 * 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.impl.verification;

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

import javax.xml.namespace.QName;
import javax.xml.transform.dom.DOMSource;
import javax.xml.validation.Schema;
import javax.xml.validation.Validator;

import at.gv.egiz.eaaf.core.exceptions.EaafProtocolException;
import at.gv.egiz.eaaf.core.exceptions.InvalidProtocolRequestException;
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.api.metadata.IRefreshableMetadataProvider;
import at.gv.egiz.eaaf.modules.pvp2.exception.SamlAssertionValidationExeption;
import at.gv.egiz.eaaf.modules.pvp2.exception.SchemaValidationException;
import at.gv.egiz.eaaf.modules.pvp2.impl.message.InboundMessage;
import at.gv.egiz.eaaf.modules.pvp2.impl.message.PvpSProfileRequest;
import at.gv.egiz.eaaf.modules.pvp2.impl.message.PvpSProfileResponse;

import org.apache.commons.lang3.StringUtils;
import org.joda.time.DateTime;
import org.opensaml.core.criterion.EntityIdCriterion;
import org.opensaml.saml.common.xml.SAMLConstants;
import org.opensaml.saml.common.xml.SAMLSchemaBuilder;
import org.opensaml.saml.common.xml.SAMLSchemaBuilder.SAML1Version;
import org.opensaml.saml.criterion.EntityRoleCriterion;
import org.opensaml.saml.criterion.ProtocolCriterion;
import org.opensaml.saml.saml2.core.Assertion;
import org.opensaml.saml.saml2.core.Audience;
import org.opensaml.saml.saml2.core.AudienceRestriction;
import org.opensaml.saml.saml2.core.Conditions;
import org.opensaml.saml.saml2.core.EncryptedAssertion;
import org.opensaml.saml.saml2.core.RequestAbstractType;
import org.opensaml.saml.saml2.core.Response;
import org.opensaml.saml.saml2.core.StatusCode;
import org.opensaml.saml.saml2.core.StatusResponseType;
import org.opensaml.saml.saml2.encryption.Decrypter;
import org.opensaml.saml.saml2.encryption.EncryptedElementTypeEncryptedKeyResolver;
import org.opensaml.saml.saml2.metadata.IDPSSODescriptor;
import org.opensaml.saml.saml2.metadata.SPSSODescriptor;
import org.opensaml.saml.security.impl.SAMLSignatureProfileValidator;
import org.opensaml.security.credential.UsageType;
import org.opensaml.security.criteria.UsageCriterion;
import org.opensaml.xmlsec.encryption.support.ChainingEncryptedKeyResolver;
import org.opensaml.xmlsec.encryption.support.DecryptionException;
import org.opensaml.xmlsec.encryption.support.EncryptedKeyResolver;
import org.opensaml.xmlsec.encryption.support.InlineEncryptedKeyResolver;
import org.opensaml.xmlsec.encryption.support.SimpleRetrievalMethodEncryptedKeyResolver;
import org.opensaml.xmlsec.keyinfo.impl.StaticKeyInfoCredentialResolver;
import org.opensaml.xmlsec.signature.support.SignatureException;
import org.opensaml.xmlsec.signature.support.SignatureTrustEngine;
import org.springframework.beans.factory.annotation.Autowired;
import org.w3c.dom.Element;
import org.xml.sax.SAXException;

import lombok.extern.slf4j.Slf4j;
import net.shibboleth.utilities.java.support.net.BasicURLComparator;
import net.shibboleth.utilities.java.support.net.URIException;
import net.shibboleth.utilities.java.support.resolver.CriteriaSet;

@Slf4j
public class SamlVerificationEngine {
  private static SAMLSchemaBuilder schemaBuilder = new SAMLSchemaBuilder(SAML1Version.SAML_11);

  private static final String ERROR_03 = "internal.pvp.03";
  private static final String ERROR_10 = "internal.pvp.10";
  private static final String ERROR_14 = "internal.pvp.14";
  private static final String ERROR_15 = "internal.pvp.15";
  private static final String ERROR_16 = "internal.pvp.16";
  private static final String ERROR_17 = "internal.pvp.17";

  private static final Object SIG_VAL_ERROR_MSG = "Signature verification return false";

  /**
   * 5 allow 3 minutes time jitter in before validation.
   */
  private static final int TIME_JITTER = 3;





  @Autowired(required = true)
  IPvp2MetadataProvider metadataProvider;

  /**
   * Verify signature of a signed SAML2 object.
   *
   *<p>This method only perform signature verification</p>
   *
   * @param msg            SAML2 message
   * @param sigTrustEngine TrustEngine
   * @throws org.opensaml.xml.security.SecurityException In case of invalid
   *                                                     signature
   * @throws Exception                                   In case of a general
   *                                                     error
   */
  public void verify(final InboundMessage msg, final SignatureTrustEngine sigTrustEngine)
      throws SecurityException, Exception {
    try {
      if (msg instanceof PvpSProfileRequest
          && ((PvpSProfileRequest) msg).getSamlRequest() instanceof RequestAbstractType) {
        verifyRequest((RequestAbstractType) ((PvpSProfileRequest) msg).getSamlRequest(),
            sigTrustEngine);
      } else if (msg instanceof PvpSProfileResponse) {
        verifyIdpResponse(((PvpSProfileResponse) msg).getResponse(), sigTrustEngine);

      } else {
        log.warn("SAML2 message type: {} not supported", msg.getClass().getName());
        throw new EaafProtocolException("internal.pvp.99", null);

      }

    } catch (final InvalidProtocolRequestException e) {
      if (StringUtils.isEmpty(msg.getEntityID())) {
        throw e;

      }
      log.debug("PVP2X message validation FAILED. Relead metadata for entityID: {}",
          msg.getEntityID());

      if (metadataProvider == null || !(metadataProvider instanceof IRefreshableMetadataProvider)
          || !((IRefreshableMetadataProvider) metadataProvider)
              .refreshMetadataProvider(msg.getEntityID())) {
        throw e;

      } else {
        log.trace("PVP2X metadata reload finished. Check validate message again.");

        if (msg instanceof PvpSProfileRequest
            && ((PvpSProfileRequest) msg).getSamlRequest() instanceof RequestAbstractType) {
          verifyRequest((RequestAbstractType) ((PvpSProfileRequest) msg).getSamlRequest(),
              sigTrustEngine);

        } else {
          verifyIdpResponse(((PvpSProfileResponse) msg).getResponse(), sigTrustEngine);

        }

      }
      log.trace("Second PVP2X message validation finished");

    }
  }

  /**
   * Verify the signature of a signed SAML2 object from ServiceProvider.
   *
   * @param samlObj signed Response from ServiceProvider
   * @param sigTrustEngine TrustEngie for verification
   * @throws InvalidProtocolRequestException In case of a verification error
   */
  public void verifySloResponse(final StatusResponseType samlObj,
      final SignatureTrustEngine sigTrustEngine) throws InvalidProtocolRequestException {
    verifyResponse(samlObj, sigTrustEngine, SPSSODescriptor.DEFAULT_ELEMENT_NAME);

  }

  /**
   * Verify the signature of a signed SAML2 object from IDP.
   *
   * <p>This method only perform signature verification</p>
   *
   * @param samlObj signed SAML2 message from IDP
   * @param sigTrustEngine TrustEngie for verification
   * @throws InvalidProtocolRequestException In case of a verification error
   */
  public void verifyIdpResponse(final StatusResponseType samlObj,
      final SignatureTrustEngine sigTrustEngine) throws InvalidProtocolRequestException {
    verifyResponse(samlObj, sigTrustEngine, IDPSSODescriptor.DEFAULT_ELEMENT_NAME);

  }

  /**
   * Validate a PVP response and all included assertions.
   *
   *<p>
   * If the SAML2 assertions are encrypted than they will be decrypted afterwards
   * </p>
   *
   *<p>This method <b>DOES NOT</b> verify the <i>Destination</i> attribute in SAML2 Response</p>
   *
   * @param samlResp            SAML2 Response object
   * @param assertionDecryption Assertion decryption-credentials to decrypt SAML2
   *                            assertions
   * @param spEntityID          EntityId of the SAML2 client
   * @param loggerName          Name for logging purposes
   * @throws SamlAssertionValidationExeption In case of a validation error
   */
  public void validateAssertion(Response samlResp,
      EaafX509Credential assertionDecryption, String spEntityID, String loggerName)
      throws SamlAssertionValidationExeption {
    validateAssertion(samlResp, assertionDecryption, spEntityID, loggerName, true);

  }

  /**
   * Validate each SAML2 assertions in a SAML2 response. <br>
   * <p>
   * If the SAML2 assertions are encrypted than they will be decrypted afterwards
   * </p>
   *
   *<p>This method <b>DOES NOT</b> verify the <i>Destination</i> attribute in SAML2 Response</p>
   *
   * @param samlResp            SAML2 Response object
   * @param assertionDecryption Assertion decryption-credentials to decrypt SAML2
   *                            assertions
   * @param spEntityID          EntityId of the SAML2 client
   * @param loggerName          Name for logging purposes
   * @param validateDateTime    <code>true</code> if <i>getIssueInstant</i>
   *                            attribute should be validated, otherwise false
   * @throws SamlAssertionValidationExeption In case of a validation error
   */
  public void validateAssertion(Response samlResp, EaafX509Credential assertionDecryption,
      String spEntityID, String loggerName, boolean validateDateTime)
      throws SamlAssertionValidationExeption {
    try {
      // pre-validate the SAML2 response
      assertionPreValidation(samlResp, loggerName, validateDateTime);

      // get Assertion from response and decrypt them if the are encrypted
      final List<Assertion> saml2assertions = getOrDecryptAndGetAssertions(samlResp, assertionDecryption);

      // validate each assertion
      final List<Assertion> validatedassertions = new ArrayList<>();
      for (final Assertion saml2assertion : saml2assertions) {
        if (internalAssertionValidation(saml2assertion, spEntityID, validateDateTime)) {
          log.debug("Add valid Assertion:" + saml2assertion.getID());
          validatedassertions.add(saml2assertion);

        } else {
          log.warn("Remove non-valid Assertion:" + saml2assertion.getID());
        }

      }

      if (validatedassertions.isEmpty()) {
        log.info("No valid PVP 2.1 assertion received.");
        throw new SamlAssertionValidationExeption(ERROR_15, new Object[] { loggerName });

      }

      samlResp.getAssertions().clear();
      samlResp.getEncryptedAssertions().clear();
      samlResp.getAssertions().addAll(validatedassertions);

    } catch (final DecryptionException e) {
      log.warn("Assertion decrypt FAILED.", e);
      throw new SamlAssertionValidationExeption(ERROR_16,
          new Object[] { e.getMessage() }, e);

//    } catch (final ConfigurationException e) {
//      throw new AssertionValidationExeption("pvp.12",
//          new Object[]{loggerName, e.getMessage()}, e);
    }
  }

  private boolean internalAssertionValidation(Assertion saml2assertion, String spEntityId,
      boolean validateDateTime) {
    boolean isAssertionValid = true;
    try {
      // schema validation
      performSchemaValidation(saml2assertion.getDOM());

      // validate DateTime conditions
      final Conditions conditions = saml2assertion.getConditions();
      if (conditions != null) {
        final DateTime notbefore = conditions.getNotBefore().minusMinutes(5);
        final DateTime notafter = conditions.getNotOnOrAfter();
        if (validateDateTime
            && (notbefore.isAfterNow() || notafter.isBeforeNow())) {
          isAssertionValid = false;
          log.info("Assertion with ID:{} is out of Date. [ Current:{}  NotBefore:{} NotAfter:{} ]",
              saml2assertion.getID(), new DateTime(), notbefore, notafter);

        }

        // validate audienceRestrictions are valid for this SP
        final List<AudienceRestriction> audienceRest = conditions.getAudienceRestrictions();
        if (audienceRest == null || audienceRest.size() == 0) {
          log.info("Assertion with ID:{} has not 'AudienceRestriction' element",
              saml2assertion.getID());
          isAssertionValid = false;

        } else {
          for (final AudienceRestriction el : audienceRest) {
            for (final Audience audience : el.getAudiences()) {
              if (!urlCompare(spEntityId, audience.getAudienceURI())) {
                log.info("Assertion with ID:{} 'AudienceRestriction' is not valid.",
                    saml2assertion.getID());
                isAssertionValid = false;

              }
            }
          }
        }

      } else {
        log.info("Assertion with ID:{} contains not 'Conditions' element",
            saml2assertion.getID());
        isAssertionValid = false;

      }

    } catch (final SchemaValidationException e) {
      isAssertionValid = false;
      log.info("Assertion with ID:{} FAILED Schema validation. Msg: {}",
          saml2assertion.getID(), e.getMessage());

    } catch (final URIException e) {
      isAssertionValid = false;
      log.info("Assertion with ID:{} FAILED AudienceRestriction validation. Msg:",
          saml2assertion.getID(), e.getMessage());

    }

    return isAssertionValid;

  }

  private List<Assertion> getOrDecryptAndGetAssertions(Response samlResp,
      EaafX509Credential assertionDecryption) throws DecryptionException {
    final List<Assertion> saml2assertions = new ArrayList<>();

    // check encrypted Assertions
    final List<EncryptedAssertion> encryAssertionList = samlResp.getEncryptedAssertions();
    if (encryAssertionList != null && encryAssertionList.size() > 0) {
      // decrypt assertions
      log.debug("Found encryped assertion. Start decryption ...");
      final List<EncryptedKeyResolver> listOfKeyResolvers = new ArrayList<>();
      listOfKeyResolvers.add(new InlineEncryptedKeyResolver());
      listOfKeyResolvers.add(new EncryptedElementTypeEncryptedKeyResolver());
      listOfKeyResolvers.add(new SimpleRetrievalMethodEncryptedKeyResolver());

      final Decrypter samlDecrypter = new Decrypter(null,
          new StaticKeyInfoCredentialResolver(assertionDecryption),
          new ChainingEncryptedKeyResolver(listOfKeyResolvers));

      for (final EncryptedAssertion encAssertion : encryAssertionList) {
        saml2assertions.add(samlDecrypter.decrypt(encAssertion));

      }
      log.debug("Assertion decryption finished. ");

    }

    saml2assertions.addAll(samlResp.getAssertions());

    return saml2assertions;

  }

  private void performSchemaValidation(final Element source) throws SchemaValidationException {

    String err = null;
    try {
      final Schema test = schemaBuilder.getSAMLSchema();
      final Validator val = test.newValidator();
      val.validate(new DOMSource(source));
      log.debug("Schema validation check done OK");
      return;

    } catch (final SAXException e) {
      err = e.getMessage();
      if (log.isDebugEnabled() || log.isTraceEnabled()) {
        log.warn("Schema validation FAILED with exception:", e);
      } else {
        log.warn("Schema validation FAILED with message: " + e.getMessage());
      }

    } catch (final Exception e) {
      err = e.getMessage();
      if (log.isDebugEnabled() || log.isTraceEnabled()) {
        log.warn("Schema validation FAILED with exception:", e);
      } else {
        log.warn("Schema validation FAILED with message: " + e.getMessage());
      }

    }

    throw new SchemaValidationException(ERROR_03, new Object[] { err });

  }

  private void verifyResponse(final StatusResponseType samlObj,
      final SignatureTrustEngine sigTrustEngine, final QName defaultElementName)
      throws InvalidProtocolRequestException {

    final SAMLSignatureProfileValidator profileValidator = new SAMLSignatureProfileValidator();
    try {
      profileValidator.validate(samlObj.getSignature());
      performSchemaValidation(samlObj.getDOM());

    } catch (final SignatureException e) {
      log.warn("Signature is not conform to SAML signature profile", e);
      throw new InvalidProtocolRequestException(ERROR_10, new Object[] {e.getMessage() }, e);

    } catch (final SchemaValidationException e) {
      throw new InvalidProtocolRequestException(ERROR_03, new Object[] { e.getMessage() }, e);

    }

    final CriteriaSet criteriaSet = new CriteriaSet();
    criteriaSet.add(new EntityIdCriterion(samlObj.getIssuer().getValue()));
    criteriaSet.add(new ProtocolCriterion(SAMLConstants.SAML20P_NS));
    criteriaSet.add(new EntityRoleCriterion(defaultElementName));
    criteriaSet.add(new UsageCriterion(UsageType.SIGNING));

    try {
      if (!sigTrustEngine.validate(samlObj.getSignature(), criteriaSet)) {
        throw new InvalidProtocolRequestException(ERROR_10, new Object[] {SIG_VAL_ERROR_MSG});

      }

    } catch (final org.opensaml.security.SecurityException e) {
      log.warn("PVP2x message signature validation FAILED.", e);
      throw new InvalidProtocolRequestException(ERROR_10, new Object[] {e.getMessage()}, e);

    }
  }

  private void verifyRequest(final RequestAbstractType samlObj,
      final SignatureTrustEngine sigTrustEngine) throws InvalidProtocolRequestException {
    final SAMLSignatureProfileValidator profileValidator = new SAMLSignatureProfileValidator();
    try {
      profileValidator.validate(samlObj.getSignature());
      performSchemaValidation(samlObj.getDOM());

    } catch (final SignatureException e) {
      log.warn("Signature is not conform to SAML signature profile", e);
      throw new InvalidProtocolRequestException(ERROR_10, new Object[] {e.getMessage()}, e);

    } catch (final SchemaValidationException e) {
      throw new InvalidProtocolRequestException(ERROR_03, new Object[] { e.getMessage() }, e);

    }

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

    try {
      if (!sigTrustEngine.validate(samlObj.getSignature(), criteriaSet)) {
        throw new InvalidProtocolRequestException(ERROR_10, new Object[] {SIG_VAL_ERROR_MSG});

      }
    } catch (final org.opensaml.security.SecurityException e) {
      log.warn("PVP2x message signature validation FAILED.", e);
      throw new InvalidProtocolRequestException(ERROR_10, new Object[] {e.getMessage()}, e);

    }

  }

  private void assertionPreValidation(Response samlResp, String loggerName, boolean validateDateTime)
      throws SamlAssertionValidationExeption {
    if (samlResp.getStatus().getStatusCode().getValue().equals(StatusCode.SUCCESS)) {
      // validate response issueInstant
      final DateTime issueInstant = samlResp.getIssueInstant();
      if (issueInstant == null) {
        log.warn("PVP response does not include a 'IssueInstant' attribute");
        throw new SamlAssertionValidationExeption(ERROR_14,
            new Object[] { loggerName, "'IssueInstant' attribute is not included" });

      }
      if (validateDateTime && issueInstant.minusMinutes(TIME_JITTER).isAfterNow()) {
        log.warn("PVP response: IssueInstant DateTime is not valid anymore.");
        throw new SamlAssertionValidationExeption(ERROR_14,
            new Object[] { loggerName, "'IssueInstant' Time is not valid any more" });

      }

    } else {
      log.info("PVP 2.x assertion includes an error. Receive errorcode "
          + samlResp.getStatus().getStatusCode().getValue());
      throw new SamlAssertionValidationExeption(ERROR_17,
          new Object[] { loggerName,
              samlResp.getIssuer().getValue(),
              samlResp.getStatus().getStatusCode().getValue(),
              samlResp.getStatus().getStatusMessage() != null
                ? samlResp.getStatus().getStatusMessage().getMessage()
                    : " no status message" });

    }
  }

  private static boolean urlCompare(String url1, String url2) throws URIException {
    final BasicURLComparator comparator = new BasicURLComparator();
    comparator.setCaseInsensitive(false);
    return comparator.compare(url1, url2);
  }

}