/*
 * 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 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.metadata.IPvp2MetadataProvider;
import at.gv.egiz.eaaf.modules.pvp2.api.metadata.IRefreshableMetadataProvider;
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.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.RequestAbstractType;
import org.opensaml.saml.saml2.core.StatusResponseType;
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.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.resolver.CriteriaSet;

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

  @Autowired(required = true)
  IPvp2MetadataProvider metadataProvider;

  /**
   * Verify signature of a signed SAML2 object.
   *
   * @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("9999", 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");
    }
  }

  public void verifySloResponse(final StatusResponseType samlObj,
      final SignatureTrustEngine sigTrustEngine) throws InvalidProtocolRequestException {
    verifyResponse(samlObj, sigTrustEngine, SPSSODescriptor.DEFAULT_ELEMENT_NAME);

  }

  public void verifyIdpResponse(final StatusResponseType samlObj,
      final SignatureTrustEngine sigTrustEngine) throws InvalidProtocolRequestException {
    verifyResponse(samlObj, sigTrustEngine, IDPSSODescriptor.DEFAULT_ELEMENT_NAME);

  }

  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("pvp2.21", new Object[] {});

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

    }

    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("pvp2.21", new Object[] {});
      }
    } catch (final org.opensaml.security.SecurityException e) {
      log.warn("PVP2x message signature validation FAILED.", e);
      throw new InvalidProtocolRequestException("pvp2.21", new Object[] {});
    }
  }

  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("pvp2.21", new Object[] {});

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

    }

    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("pvp2.21", new Object[] {});
      }
    } catch (final org.opensaml.security.SecurityException e) {
      log.warn("PVP2x message signature validation FAILED.", e);
      throw new InvalidProtocolRequestException("pvp2.21", new Object[] {});
    }
  }

  protected 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("pvp2.22", new Object[] { err });

  }

}