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

import java.io.IOException;
import java.io.InputStream;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.ECPrivateKey;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.List;

import javax.annotation.Nonnull;
import javax.xml.namespace.QName;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.TransformerException;
import javax.xml.transform.dom.DOMSource;
import javax.xml.validation.Schema;
import javax.xml.validation.Validator;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.opensaml.core.xml.XMLObject;
import org.opensaml.core.xml.XMLObjectBuilderFactory;
import org.opensaml.core.xml.config.XMLObjectProviderRegistrySupport;
import org.opensaml.core.xml.io.Marshaller;
import org.opensaml.core.xml.io.MarshallingException;
import org.opensaml.core.xml.io.Unmarshaller;
import org.opensaml.core.xml.io.UnmarshallingException;
import org.opensaml.core.xml.schema.XSString;
import org.opensaml.core.xml.schema.impl.XSStringBuilder;
import org.opensaml.core.xml.util.XMLObjectSupport;
import org.opensaml.messaging.decoder.MessageDecodingException;
import org.opensaml.saml.common.SAMLObjectContentReference;
import org.opensaml.saml.common.xml.SAMLSchemaBuilder;
import org.opensaml.saml.common.xml.SAMLSchemaBuilder.SAML1Version;
import org.opensaml.saml.saml2.core.Attribute;
import org.opensaml.saml.saml2.core.Status;
import org.opensaml.saml.saml2.core.StatusCode;
import org.opensaml.saml.saml2.metadata.AssertionConsumerService;
import org.opensaml.saml.saml2.metadata.SPSSODescriptor;
import org.opensaml.security.SecurityException;
import org.opensaml.security.x509.X509Credential;
import org.opensaml.soap.soap11.Body;
import org.opensaml.soap.soap11.Envelope;
import org.opensaml.xmlsec.SecurityConfigurationSupport;
import org.opensaml.xmlsec.SignatureSigningConfiguration;
import org.opensaml.xmlsec.keyinfo.KeyInfoGenerator;
import org.opensaml.xmlsec.keyinfo.KeyInfoGeneratorFactory;
import org.opensaml.xmlsec.keyinfo.KeyInfoGeneratorManager;
import org.opensaml.xmlsec.keyinfo.NamedKeyInfoGeneratorManager;
import org.opensaml.xmlsec.keyinfo.impl.BasicKeyInfoGeneratorFactory;
import org.opensaml.xmlsec.signature.KeyInfo;
import org.opensaml.xmlsec.signature.SignableXMLObject;
import org.opensaml.xmlsec.signature.Signature;
import org.opensaml.xmlsec.signature.support.ContentReference;
import org.opensaml.xmlsec.signature.support.SignatureConstants;
import org.opensaml.xmlsec.signature.support.SignatureException;
import org.opensaml.xmlsec.signature.support.Signer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.xml.sax.SAXException;

import at.gv.egiz.eaaf.core.impl.utils.DomUtils;
import at.gv.egiz.eaaf.core.impl.utils.Random;
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.reqattr.EaafRequestedAttribute;
import at.gv.egiz.eaaf.modules.pvp2.exception.SamlSigningException;
import at.gv.egiz.eaaf.modules.pvp2.exception.SchemaValidationException;
import net.shibboleth.shared.xml.QNameSupport;
import net.shibboleth.shared.xml.SerializeSupport;


public class Saml2Utils {
  private static final Logger log = LoggerFactory.getLogger(Saml2Utils.class);

  private static DocumentBuilder builder;
  private static SAMLSchemaBuilder schemaBuilder;

  static {
    schemaBuilder = new SAMLSchemaBuilder(SAML1Version.SAML_11);

    final DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
    factory.setNamespaceAware(true);

    try {
      builder = factory.newDocumentBuilder();

    } catch (final ParserConfigurationException e) {
      // TODO Auto-generated catch block
      e.printStackTrace();
    }
  }

  /**
   * Sign a OpenSAML 3.x object with a {@link X509Credential}. <br>
   * <p>
   * This method used {@link PvpConstants.DEFAULT_SIGNING_METHODE_RSA} or
   * {@link PvpConstants.DEFAULT_SIGNING_METHODE_EC} as algorithm
   * </p>
   *
   * @param <T>               {@link SignableXMLObject}
   * @param toSign            object that should be signed
   * @param signingCredential Credentials that should be used for signing
   * @param injectCertificate true, if certificate should be part of the signature
   * @return Signed object
   * @throws SamlSigningException In case of a signing error
   */
  public static <T extends SignableXMLObject> T signSamlObject(@Nonnull T toSign,
      @Nonnull EaafX509Credential signingCredential, boolean injectCertificate) throws SamlSigningException {

    try {
      final String usedSigAlg = signingCredential.getSignatureAlgorithmForSigning();
      final Signature signature = createSignature(signingCredential, usedSigAlg, injectCertificate);
      toSign.setSignature(signature);

      final String digestAlgorithm = getDigestAlgorithm(usedSigAlg);
      final List<ContentReference> contentReferences = signature.getContentReferences();
      if (!CollectionUtils.isEmpty(contentReferences)) {
        ((SAMLObjectContentReference) contentReferences.get(0)).setDigestAlgorithm(digestAlgorithm);

      } else {
        log.error("Unable to set DigestMethodAlgorithm - algorithm {} not set", digestAlgorithm);

      }

      log.trace("Marshall samlToken.");
      XMLObjectProviderRegistrySupport.getMarshallerFactory().getMarshaller(toSign).marshall(toSign);

      log.trace("Sign samlToken.");
      Signer.signObject(signature);

      return toSign;

    } catch (final SignatureException | MarshallingException | SecurityException e) {
      throw new SamlSigningException("internal.pvp.96",
          new Object[] { signingCredential.getEntityId(), e.getMessage() }, e);

    }

  }

  /**
   * SAML2 message unmarshaller that performs schema validation before unmarshall
   * the message.
   *
   * @param messageStream SAML2 message that shoulld be unmarshalled
   * @return OpenSAML XML object
   * @throws MessageDecodingException In case of a schema-validation or
   *                                  unmarshalling error
   */
  public static XMLObject unmarshallMessage(final InputStream messageStream) throws MessageDecodingException {
    try {
      final Element samlElement = DomUtils.parseXmlValidating(messageStream);

      if (log.isTraceEnabled()) {
        log.trace("Resultant DOM message was:");
        log.trace(SerializeSupport.nodeToString(samlElement));
      }

      log.debug("Unmarshalling DOM parsed from InputStream");
      final Unmarshaller unmarshaller = XMLObjectSupport.getUnmarshaller(samlElement);
      if (unmarshaller == null) {
        log.error("Unable to unmarshall InputStream, no unmarshaller registered for element "
            + QNameSupport.getNodeQName(samlElement));
        throw new UnmarshallingException(
            "Unable to unmarshall InputStream, no unmarshaller registered for element "
                + QNameSupport.getNodeQName(samlElement));
      }

      final XMLObject message = unmarshaller.unmarshall(samlElement);

      log.debug("InputStream succesfully unmarshalled");

      return message;

    } catch (final UnmarshallingException e) {
      log.error("Error unmarshalling message from input stream", e);
      throw new MessageDecodingException("Error unmarshalling message from input stream", e);

    } catch (ParserConfigurationException | SAXException e) {
      log.warn("Message schema-validation failed.");
      throw new MessageDecodingException("Message schema-validation failed.",
          new SchemaValidationException("internal.pvp.03", new Object[] { e.getMessage() }, e));

    } catch (final IOException e) {
      log.error("Error read message from input stream", e);
      throw new MessageDecodingException("Error read message from input stream", e);

    }
  }

  /**
   * Select signature algorithm for a given credential.
   *
   * @param credentials     {@link X509Credential} that will be used for key operations
   * @param rsaSigAlgorithm RSA based algorithm that should be used in
   *                        case of RSA credential
   * @param ecSigAlgorithm  EC based algorithm that should be used in case
   *                        of RSA credential
   * @return either the RSA based algorithm or the EC based algorithm
   * @throws SamlSigningException In case of an unsupported credential
   */
  public static String getKeyOperationAlgorithmFromCredential(X509Credential credentials,
      String rsaSigAlgorithm, String ecSigAlgorithm) throws SamlSigningException {
    final PrivateKey privatekey = credentials.getPrivateKey();
    final PublicKey publickey = credentials.getPublicKey();
    if (privatekey instanceof RSAPrivateKey
        || publickey instanceof RSAPublicKey) {
      return rsaSigAlgorithm;

    } else if (privatekey instanceof ECPrivateKey
        || publickey instanceof ECPublicKey) {
      return ecSigAlgorithm;

    } else {
      log.warn("Could NOT evaluate the Private-Key type from " + credentials.getEntityId()
          + " credential.");
      throw new SamlSigningException("internal.pvp.97",
          new Object[] { credentials.getEntityId(),
              privatekey != null ? privatekey.getClass().getName() : publickey.getClass().getName()
              });

    }
  }

  /**
   * Select a digest algorithm for a already selected signing algorithm.
   *
   * @param signatureAlgorithmName Signing algorithm that will be used
   * @return Digest algorithm identifier
   */
  public static String getDigestAlgorithm(String signatureAlgorithmName) {
    if (StringUtils.isBlank(signatureAlgorithmName)) {
      return PvpConstants.DEFAULT_DIGESTMETHODE;
    }

    final String canonicalAlgorithm = signatureAlgorithmName.trim();
    final String digestAlgorithm = PvpConstants.SIGNATURE_TO_DIGEST_ALGORITHM_MAP.get(canonicalAlgorithm);
    if (null != digestAlgorithm) {
      return digestAlgorithm;

    }

    log.warn("Signing algorithm: {} does not contain a known digist algorithm. Use: {} as default",
        signatureAlgorithmName, PvpConstants.DEFAULT_DIGESTMETHODE);
    return PvpConstants.DEFAULT_DIGESTMETHODE;

  }

  /**
   * Get a {@link KeyInfoGenerator} that injects key information into XML
   * signature.
   *
   * @param credential        @link X509Credential} that will be used for signing
   * @param injectCertificate Set <code>true</code> if the certificate should be
   *                          added to KeyInfo
   * @return Generator for a XML signature key-information
   */
  public static KeyInfoGenerator getKeyInfoGenerator(X509Credential credential, boolean injectCertificate) {
    // OpenSAML3 only support RSA and DSA for direct key injection
    KeyInfoGeneratorFactory keyInfoGenFac = null;
    if (injectCertificate || credential.getPublicKey() instanceof ECPublicKey) {
      final SignatureSigningConfiguration secConfiguration = SecurityConfigurationSupport
          .getGlobalSignatureSigningConfiguration();
      final NamedKeyInfoGeneratorManager keyInfoManager = secConfiguration.getKeyInfoGeneratorManager();
      final KeyInfoGeneratorManager keyInfoGenManager = keyInfoManager.getDefaultManager();
      keyInfoGenFac = keyInfoGenManager.getFactory(credential);

    } else {
      keyInfoGenFac = createKeyInfoWithoutCertificate();

    }

    return keyInfoGenFac.newInstance();

  }

  /**
   * Create a SAML2 object.
   *
   * @param <T>   SAML2 object class
   * @param clazz object class
   * @return SAML2 object
   */
  public static <T> T createSamlObject(final Class<T> clazz) {
    try {
      final XMLObjectBuilderFactory builderFactory =
          XMLObjectProviderRegistrySupport.getBuilderFactory();

      final QName defaultElementName =
          (QName) clazz.getDeclaredField("DEFAULT_ELEMENT_NAME").get(null);
      @SuppressWarnings("unchecked")
      final T object =
          (T) builderFactory.getBuilder(defaultElementName).buildObject(defaultElementName);
      return object;
    } catch (final Throwable e) {
      e.printStackTrace();
      return null;
    }
  }

  /**
   * Get a new SAML2 conform random value.
   *
   * @return
   */
  public static String getSecureIdentifier() {
    return "_".concat(Random.nextHexRandom16());

  }

  /**
   * Transform SAML2 Object to Element.
   *
   * @param object SAML2 object
   * @return Element
   * @throws IOException          In case of an transformation error
   * @throws MarshallingException In case of an transformation error
   * @throws TransformerException In case of an transformation error
   */
  public static Document asDomDocument(final XMLObject object)
      throws IOException, MarshallingException, TransformerException {
    final Document document = builder.newDocument();
    final Marshaller out =
        XMLObjectProviderRegistrySupport.getMarshallerFactory().getMarshaller(object);
    out.marshall(object, document);
    return document;
  }

  /**
   * Build success status element.
   *
   * @return
   */
  public static Status getSuccessStatus() {
    final Status status = Saml2Utils.createSamlObject(Status.class);
    final StatusCode statusCode = Saml2Utils.createSamlObject(StatusCode.class);
    statusCode.setValue(StatusCode.SUCCESS);
    status.setStatusCode(statusCode);
    return status;
  }

  /**
   * Get AssertionConsumerService Index from metadata element.
   *
   * @param spSsoDescriptor metadata element
   * @return
   */
  public static int getDefaultAssertionConsumerServiceIndex(final SPSSODescriptor spSsoDescriptor) {

    final List<AssertionConsumerService> assertionConsumerList =
        spSsoDescriptor.getAssertionConsumerServices();

    for (final AssertionConsumerService el : assertionConsumerList) {
      if (el.isDefault()) {
        return el.getIndex();
      }

    }

    return 0;
  }

  /**
   * Build SOAP11 body from SAML2 object.
   *
   * @param payload SAML2 object
   * @return
   */
  public static Envelope buildSoap11Envelope(final XMLObject payload) {
    final XMLObjectBuilderFactory bf = XMLObjectProviderRegistrySupport.getBuilderFactory();
    final Envelope envelope = (Envelope) bf.getBuilder(Envelope.DEFAULT_ELEMENT_NAME)
        .buildObject(Envelope.DEFAULT_ELEMENT_NAME);
    final Body body =
        (Body) bf.getBuilder(Body.DEFAULT_ELEMENT_NAME).buildObject(Body.DEFAULT_ELEMENT_NAME);

    body.getUnknownXMLObjects().add(payload);
    envelope.setBody(body);

    return envelope;
  }

  /**
   * Generate EAAF specific requested attribute.
   *
   * @param attr       SAML2 attribute definition
   * @param isRequired is-mandatory flag
   * @param value      Attribute value
   * @return
   */
  public static EaafRequestedAttribute generateReqAuthnAttributeSimple(
      final Attribute attr, final boolean isRequired, final String value) {
    final EaafRequestedAttribute requested =
        Saml2Utils.createSamlObject(EaafRequestedAttribute.class);
    requested.setName(attr.getName());
    requested.setNameFormat(attr.getNameFormat());
    requested.setFriendlyName(attr.getFriendlyName());
    requested.setIsRequired(String.valueOf(isRequired));
    final List<XMLObject> attributeValues = requested.getAttributeValues();
    if (StringUtils.isNotEmpty(value)) {
      final XMLObject attributeValueForRequest =
          createAttributeValue(PvpConstants.EIDAS_REQUESTED_ATTRIBUTE_VALUE_TYPE, value);
      attributeValues.add(attributeValueForRequest);
    }
    return requested;

  }

  /**
   * Perform XML schema-validation on SAML2 object.
   *
   * @param xmlObject SAML2 object
   * @throws Exception In case of a validation error
   */
  public static void schemeValidation(final XMLObject xmlObject) throws Exception {
    try {

      final Schema test = schemaBuilder.getSAMLSchema();
      final Validator val = test.newValidator();
      final DOMSource source = new DOMSource(xmlObject.getDOM());
      val.validate(source);
      log.debug("SAML2 Scheme validation successful");
      return;

    } catch (final Exception e) {
      log.warn("SAML2 scheme validation FAILED.", e);
      throw e;

    }
  }

  private static XMLObject createAttributeValue(final QName attributeValueType,
      final String value) {
    final XSStringBuilder stringBuilder = (XSStringBuilder) XMLObjectProviderRegistrySupport
        .getBuilderFactory().getBuilder(XSString.TYPE_NAME);
    final XSString stringValue = stringBuilder.buildObject(attributeValueType, XSString.TYPE_NAME);
    stringValue.setValue(value);
    return stringValue;

  }

  private static Signature createSignature(X509Credential signingCredential,
      String usedSigAlg, boolean injectCertificate)
      throws SecurityException, SamlSigningException {
    log.trace("Generating OpenSAML signature object ... ");
    final Signature signature = (Signature) XMLObjectProviderRegistrySupport.getBuilderFactory()
        .getBuilder(Signature.DEFAULT_ELEMENT_NAME)
        .buildObject(Signature.DEFAULT_ELEMENT_NAME);
    signature.setSigningCredential(signingCredential);
    signature.setSignatureAlgorithm(usedSigAlg);
    final KeyInfo keyInfo = getKeyInfoGenerator(signingCredential, injectCertificate).generate(
        signingCredential);
    signature.setKeyInfo(keyInfo);
    signature.setCanonicalizationAlgorithm(SignatureConstants.ALGO_ID_C14N_EXCL_OMIT_COMMENTS);
    return signature;

  }

  private static KeyInfoGeneratorFactory createKeyInfoWithoutCertificate() {
    final KeyInfoGeneratorFactory keyInfoGenFac = new BasicKeyInfoGeneratorFactory();
    ((BasicKeyInfoGeneratorFactory) keyInfoGenFac).setEmitPublicKeyValue(true);
    ((BasicKeyInfoGeneratorFactory) keyInfoGenFac).setEmitEntityIDAsKeyName(true);
    ((BasicKeyInfoGeneratorFactory) keyInfoGenFac).setEmitKeyNames(true);
    ((BasicKeyInfoGeneratorFactory) keyInfoGenFac).setEmitPublicDEREncodedKeyValue(true);
    return keyInfoGenFac;
  }

}