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

import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;

import javax.xml.XMLConstants;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.w3c.dom.Document;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;

/**
 * Utilities for handling data in DOM format.
 *
 */
public final class DomUtils {

    private static final String DOM_READ_ERROR_MESSAGE =
        "Could not read DOM data from input stream";

    private static final String DOM_WRITE_HEADER_ERROR_MESSAGE =
        "Could not write DOM header to output stream";

    private static final String DOM_DATA_WRITE_ERROR_MESSAGE =
        "Could not write DOM data to output stream";

    private static final String EXTERNAL_GENERAL_ENTITIES_ALLOW_FEATURE_URL =
        "http://xml.org/sax/features/external-general-entities";

    private static final String EXTERNAL_PARAMETER_ENTITIES_ALLOW_FEATURE_URL =
        "http://xml.org/sax/features/external-parameter-entities";

    /**
     * This class cannot be instantiated.
     */
    private DomUtils() {

    }

    /**
     * Reads some DOM data from a ${@link java.io.InputStream} and returns it as
     * a ${@link org.w3c.dom.Document}.
     *
     * @param inStream
     *            the ${@link java.io.InputStream} from which to read the DOM
     *            data.
     * @return the ${@link org.w3c.dom.Document} containing the DOM data.
     * @throws GeneralSecurityException
     *             if any IO or parsing error occurs.
     */
    public static Document read(final InputStream inStream)
            throws GeneralSecurityException {

        // Create DOM builder instance.
        DocumentBuilder domBuilder;
        try {
            DocumentBuilderFactory xmlBuilderFactory =
                DocumentBuilderFactory.newInstance();
            xmlBuilderFactory.setNamespaceAware(true);
            xmlBuilderFactory.setFeature(
                EXTERNAL_GENERAL_ENTITIES_ALLOW_FEATURE_URL, false);
            xmlBuilderFactory.setFeature(
                EXTERNAL_PARAMETER_ENTITIES_ALLOW_FEATURE_URL, false);
            domBuilder = xmlBuilderFactory.newDocumentBuilder();
        } catch (ParserConfigurationException e) {
            throw new GeneralSecurityException(
                "Could not create DOM builder instance.", e);
        }

        // Read in DOM data and write to ${@link org.w3c.com.Document}
        // object.
        try {
            return domBuilder.parse(new InputSource(inStream));
        } catch (IOException e) {
            throw new GeneralSecurityException(DOM_READ_ERROR_MESSAGE, e);
        } catch (SAXException e) {
            throw new GeneralSecurityException(DOM_READ_ERROR_MESSAGE, e);
        } finally {
            closeQuietly(inStream);
        }
    }

    /**
     * Writes the DOM data of a ${@link org.w3c.dom.Document} to a
     * {@link java.io.OutputStream}.
     *
     * @param dom
     *            the ${@link org.w3c.dom.Document} containing the DOM data.
     * @param outStream
     *            the ${@link java.io.OutputStream} to which to write the DOM
     *            data.
     * @throws GeneralSecurityException
     *             if IO or others error occur.
     */
    public static void write(final Document dom,
            final OutputStream outStream) throws GeneralSecurityException {

        // Set DOM encoding.
        String encoding = dom.getInputEncoding();
        if (encoding == null) {
            encoding = StandardCharsets.UTF_8.name();
        }

        // Construct DOM header and write to output stream.
        try {
            StringBuilder header = new StringBuilder("<?xml version=\"").append(dom.getXmlVersion()).append("\"")
                .append(" encoding=\"").append(encoding);
            if (dom.getXmlStandalone()) {
                header.append("\" standalone=\"no\"");
            }
            header.append("\"?>\n");
            outStream.write(header.toString().getBytes(
                StandardCharsets.UTF_8));
        } catch (IOException e) {
            throw new GeneralSecurityException(
                DOM_WRITE_HEADER_ERROR_MESSAGE, e);
        }

        // Write DOM data to output stream.
        try {
            TransformerFactory transformerFactory =
                TransformerFactory.newInstance();
            transformerFactory.setAttribute(
                XMLConstants.ACCESS_EXTERNAL_DTD, "");

            Transformer transformer = transformerFactory.newTransformer();
            transformer.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION,
                "yes");
            transformer.transform(new DOMSource(dom), new StreamResult(
                outStream));
        } catch (TransformerException e) {
            throw new GeneralSecurityException(
                DOM_DATA_WRITE_ERROR_MESSAGE, e);
        } finally {
            closeQuietly(outStream);
        }
    }

    /**
     * Retrieves a unique {@link Node} from some DOM data. An exception is
     * thrown if the {@link Node} cannot be found or is not unique.
     *
     * @param dom
     *            the ${@link org.w3c.dom.Document} containing the DOM data.
     * @param nodeName
     *            the name of the {@link Node}.
     * @return the unique {@link Node}.
     * @throws GeneralSecurityException
     *             if the Node is not unique.
     */
    public static Node getUniqueNode(final Document dom,
            final String nodeName) throws GeneralSecurityException {

        NodeList nodeList = dom.getElementsByTagName(nodeName);
        if (nodeList.getLength() != 1) {
            throw new GeneralSecurityException(
                "Expected to find exactly one node " + nodeName
                    + ". Found " + nodeList.getLength());
        }

        return nodeList.item(0);
    }

    /**
     * Retrieves a unique {@link Node} from some DOM data, with a specified
     * parent node. An exception is thrown if the {@link Node} cannot be found,
     * is not unique or does not have the specified parent node.
     *
     * @param dom
     *            the ${@link org.w3c.dom.Document} containing the DOM data.
     * @param nodeName
     *            the name of the {@link Node}.
     * @param parentNodeName
     *            the name of the parent {@link Node}.
     * @return the unique {@link Node}.
     * @throws GeneralSecurityException
     *             if node is not unique or the node has unexpected parent.
     */
    public static Node getUniqueNode(final Document dom,
            final String nodeName, final String parentNodeName)
            throws GeneralSecurityException {

        Node foundNode = getUniqueNode(dom, nodeName);

        String foundParentNodeName =
            foundNode.getParentNode().getNodeName();
        if (!foundParentNodeName.equals(parentNodeName)) {
            throw new GeneralSecurityException("The parent node of node "
                + nodeName + " was expected to be " + parentNodeName
                + " but it is " + foundParentNodeName);
        }

        return foundNode;
    }
    
    private static void closeQuietly(Closeable closeable) {
        try {
            closeable.close();
        } catch (IOException e) {
            // nothing to do
        }
    }
}
