/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.cryptolib.asymmetric.xml.signer.factory;

import java.io.InputStream;
import java.io.OutputStream;
import java.security.GeneralSecurityException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.Certificate;
import java.util.Arrays;
import java.util.Collections;

import javax.xml.crypto.MarshalException;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureException;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.dom.DOMValidateContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;

import org.w3c.dom.Document;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.asymmetric.utils.DomUtils;
import com.scytl.cryptolib.asymmetric.xml.signer.configuration.XmlDigitalSignerPolicy;
import com.scytl.cryptolib.commons.configuration.Provider;

/**
 * A digital signer for data in XML format.
 * <P>
 * Instances of this class are immutable.
 */
public class CryptoXmlDigitalSigner {

    /**
     * The node name of the digital signature of signed data in DOM format.
     */
    public static final String DIGITAL_SIGNATURE_DOM_NODE_NAME =
        "Signature";

    private static final String DIGITAL_SIGNER_FACTORY_CONFIGURATION_ERROR_MESSAGE =
        "Could not configure XML digital signer factory.";

    private static final String DIGITAL_SIGNER_FACTORY_REFERENCE_CREATION_ERROR_MESSAGE =
        "Could not create XML digital signer factory reference.";

    private static final String DIGITAL_SIGNER_SIGNED_INFO_ELEMENT_CREATION_ERROR_MESSAGE =
        "Could not create XML digital signer signed info element.";

    private static final String XML_DOCUMENT_TO_SIGN_READ_ERROR_MESSAGE =
        "Could not read XML document to sign.";

    private static final String DIGITAL_SIGNER_CONTEXT_CREATION_ERROR_MESSAGE =
        "Could not create XML signer context.";

    private static final String DIGITAL_SIGNING_ERROR_MESSAGE =
        "Could not digitally sign the XML content.";

    private static final String SIGNED_XML_DOCUMENT_WRITE_ERROR_MESSAGE =
        "Could not write signed XML document.";

    private static final String SIGNED_XML_DOCUMENT_READ_ERROR_MESSAGE =
        "Could not read signed XML document.";

    private static final String DIGITAL_SIGNATURE_VERIFIER_CONTEXT_CREATION_ERROR_MESSAGE =
        "Could not create XML signature verifier context.";

    private static final String DIGITAL_SIGNATURE_VERIFIER_CREATION_ERROR_MESSAGE =
        "Could not create XML signature verifier.";

    private static final String DIGITAL_SIGNER_VERIFICATION_ERROR_MESSAGE =
        "Could not verify XML data signature.";

    private final String _digitalSignatureAlgorithm;

    private final String _messageDigestAlgorithm;

    private final String _transformMethod;

    private final String _canonicalizationMethod;

    private XMLSignatureFactory _xmlSignerFactory;

    private final KeyInfoFactory _keyInfoFactory;

    /**
     * Creates an instance of signer with the provided policy.
     *
     * @param xmlDigitalSignerPolicy
     *            policy for XML digital signers.
     * @throws CryptoLibException
     *             if creation of the signer fails.
     */
    CryptoXmlDigitalSigner(
            final XmlDigitalSignerPolicy xmlDigitalSignerPolicy) {

        // Retrieve XML digital signer properties.
        _digitalSignatureAlgorithm =
            xmlDigitalSignerPolicy.getXmlDigitalSignerAlgorithmAndSpec()
                .getDigitalSignatureAlgorithm();
        _messageDigestAlgorithm =
            xmlDigitalSignerPolicy.getXmlDigitalSignerAlgorithmAndSpec()
                .getMessageDigestAlgorithm();
        _transformMethod =
            xmlDigitalSignerPolicy.getXmlDigitalSignerAlgorithmAndSpec()
                .getTransformMethod();
        _canonicalizationMethod =
            xmlDigitalSignerPolicy.getXmlDigitalSignerAlgorithmAndSpec()
                .getCanonicalizationMethod();
        String mechanismType =
            xmlDigitalSignerPolicy.getXmlDigitalSignerAlgorithmAndSpec()
                .getMechanismType();
        String provider =
            xmlDigitalSignerPolicy.getXmlDigitalSignerAlgorithmAndSpec()
                .getProvider().getProviderName();

        // Create XML digital signer factory instance, based on provider.
        try {
            if (Provider.DEFAULT == xmlDigitalSignerPolicy
                .getXmlDigitalSignerAlgorithmAndSpec().getProvider()) {
                _xmlSignerFactory =
                    XMLSignatureFactory.getInstance(mechanismType);
            } else {
                _xmlSignerFactory =
                    XMLSignatureFactory.getInstance(mechanismType,
                        provider);
            }
        } catch (GeneralSecurityException e) {
            throw new CryptoLibException(
                "Failed to create XML digital signer factory in this environment. Attempted to use the provider: "
                    + provider
                    + ", and the XML mechanism type: "
                    + mechanismType
                    + ". Error message was "
                    + e.getMessage(), e);
        }

        // Configure XML digital signer factory.
        try {
            _xmlSignerFactory.newSignatureMethod(
                _digitalSignatureAlgorithm, null);
            _xmlSignerFactory.newDigestMethod(_messageDigestAlgorithm,
                null);
            _xmlSignerFactory.newTransform(_transformMethod,
                (TransformParameterSpec) null);
            _xmlSignerFactory.newCanonicalizationMethod(
                _canonicalizationMethod, (C14NMethodParameterSpec) null);
        } catch (GeneralSecurityException e) {
            throw new CryptoLibException(
                DIGITAL_SIGNER_FACTORY_CONFIGURATION_ERROR_MESSAGE, e);
        }

        _keyInfoFactory = KeyInfoFactory.getInstance(mechanismType);
    }

    /**
     * Digitally signs some data that is in XML format.
     *
     * @param key
     *            the {@link java.security.PrivateKey} to sign the XML data.
     * @param certChain
     *            the {@link java.security.cert.Certificate} chain (optional).
     * @param inStream
     *            the ${@link java.io.InputStream} from which to read the XML
     *            data to be signed.
     * @param outStream
     *            the ${@link java.io.OutputStream} in which to write the signed
     *            XML data.
     * @param signatureParentNode
     *            the name of the parent node of the signature node in the
     *            signed XML data.
     * @throws GeneralCryptoLibException
     *             if the given private key is invalid, the XML data cannot be
     *             read, the signature generation process fails or the signed
     *             XML data cannot be written.
     */
    public void sign(final PrivateKey key, final Certificate[] certChain,
            final InputStream inStream, final OutputStream outStream,
            final String signatureParentNode)
            throws GeneralCryptoLibException {

        Document dom = readXMLToDOM(inStream);

        DOMSignContext xmlSignerContext =
            setupXMLDigitalSignerContext(key, signatureParentNode, dom);

        SignedInfo signedInfo = constructSignedInfo();

        KeyInfo keyInfo = constructKeyInfo(certChain);

        signXMLData(xmlSignerContext, signedInfo, keyInfo);

        writeXMLDataToOutputStream(outStream, dom);
    }

    /**
     * Verifies the signature of some data that is in XML format.
     *
     * @param key
     *            the {@link java.security.PublicKey} to verify the signature of
     *            the XML data.
     * @param inStream
     *            the ${@link java.io.InputStream} from which to read the signed
     *            XML data.
     * @param signatureParentNode
     *            name of parent node of signature node in signed XML data.
     * @return true if verification was completed successfully and false
     *         otherwise.
     * @throws GeneralCryptoLibException
     *             if the given public key is invalid, the signed XML data
     *             cannot be read or the signature verification process fails.
     */
    public boolean verify(final PublicKey key, final InputStream inStream,
            final String signatureParentNode)
            throws GeneralCryptoLibException {

        // Read signed XML data into DOM object.
        Document dom;
        try {
            dom = DomUtils.read(inStream);
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                SIGNED_XML_DOCUMENT_READ_ERROR_MESSAGE, e);
        }

        // Setup XML digital signature verifier context.
        DOMValidateContext xmlSignatureVerifierContext;
        try {
            xmlSignatureVerifierContext =
                new DOMValidateContext(key, DomUtils.getUniqueNode(dom,
                    DIGITAL_SIGNATURE_DOM_NODE_NAME, signatureParentNode));
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                DIGITAL_SIGNATURE_VERIFIER_CONTEXT_CREATION_ERROR_MESSAGE,
                e);
        }
        xmlSignatureVerifierContext.setBaseURI("");

        // Create XML digital signature verifier instance.
        XMLSignature xmlSignatureVerifier;
        try {
            xmlSignatureVerifier =
                _xmlSignerFactory
                    .unmarshalXMLSignature(xmlSignatureVerifierContext);
        } catch (MarshalException e) {
            throw new CryptoLibException(
                DIGITAL_SIGNATURE_VERIFIER_CREATION_ERROR_MESSAGE, e);
        }

        // Verify signature of XML data.
        try {
            return xmlSignatureVerifier
                .validate(xmlSignatureVerifierContext);
        } catch (XMLSignatureException e) {
            throw new CryptoLibException(
                DIGITAL_SIGNER_VERIFICATION_ERROR_MESSAGE, e);
        }
    }

    private DOMSignContext setupXMLDigitalSignerContext(
            final PrivateKey key, final String signatureParentNode,
            final Document dom) throws GeneralCryptoLibException {

        try {
            DOMSignContext xmlSignerContext =
                new DOMSignContext(key, DomUtils.getUniqueNode(dom,
                    signatureParentNode));
            xmlSignerContext.setBaseURI("");
            return xmlSignerContext;
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                DIGITAL_SIGNER_CONTEXT_CREATION_ERROR_MESSAGE, e);
        }
    }

    private SignedInfo constructSignedInfo()
            throws GeneralCryptoLibException {

        // Create new reference to XML digital signer factory. Use empty string
        // for first argument to specify that entire document is to be signed.
        Reference ref;
        try {
            ref =
                _xmlSignerFactory.newReference("", _xmlSignerFactory
                    .newDigestMethod(_messageDigestAlgorithm, null),
                    Collections.singletonList(_xmlSignerFactory
                        .newTransform(_transformMethod,
                            (TransformParameterSpec) null)), null, null);
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                DIGITAL_SIGNER_FACTORY_REFERENCE_CREATION_ERROR_MESSAGE, e);
        }

        // Create new SignedInfo element with XML digital signer factory.
        try {
            return _xmlSignerFactory.newSignedInfo(_xmlSignerFactory
                .newCanonicalizationMethod(_canonicalizationMethod,
                    (C14NMethodParameterSpec) null), _xmlSignerFactory
                .newSignatureMethod(_digitalSignatureAlgorithm, null),
                Collections.singletonList(ref));
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                DIGITAL_SIGNER_SIGNED_INFO_ELEMENT_CREATION_ERROR_MESSAGE,
                e);
        }
    }

    private KeyInfo constructKeyInfo(final Certificate[] certChain) {

        KeyInfo keyInfo = null;
        if (certChain != null && certChain.length > 0) {
            X509Data data =
                _keyInfoFactory.newX509Data(Arrays.asList(certChain));
            keyInfo =
                _keyInfoFactory
                    .newKeyInfo(Collections.singletonList(data));
        }

        return keyInfo;
    }

    private void signXMLData(final DOMSignContext xmlSignerContext,
            final SignedInfo signedInfo, final KeyInfo keyInfo)
            throws GeneralCryptoLibException {

        try {
            XMLSignature xmlSigner =
                _xmlSignerFactory.newXMLSignature(signedInfo, keyInfo);
            xmlSigner.sign(xmlSignerContext);
        } catch (MarshalException | XMLSignatureException e) {
            throw new CryptoLibException(DIGITAL_SIGNING_ERROR_MESSAGE, e);
        }
    }

    private void writeXMLDataToOutputStream(final OutputStream outStream,
            final Document dom) throws GeneralCryptoLibException {

        try {
            DomUtils.write(dom, outStream);
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                SIGNED_XML_DOCUMENT_WRITE_ERROR_MESSAGE, e);
        }
    }

    private Document readXMLToDOM(final InputStream inStream)
            throws GeneralCryptoLibException {

        try {
            return DomUtils.read(inStream);
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                XML_DOCUMENT_TO_SIGN_READ_ERROR_MESSAGE, e);
        }
    }
}
