package com.airlock.iam.saml2.infrastructure.plugin.configuration;

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.core.misc.util.xml.XmlTransformerFactory;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Optional;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.Result;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerException;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

public final class MetadataUtils {
   private MetadataUtils() {
   }

   @SuppressFBWarnings(value = "XXE_DOCUMENT", justification = "DocumentBuilderFactory has secure defaults and input is trusted")
   public static Document loadDocument(InputStream is) {
      try {
         DocumentBuilderFactory factory = secureDocumentBuilderFactory();
         return factory.newDocumentBuilder().parse(is);
      } catch (IOException | ParserConfigurationException | SAXException e) {
         throw new FailureException(e);
      }
   }

   public static byte[] getDocumentBytes(Document doc) {
      try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
         Result outputTarget = new StreamResult(outputStream);
         Transformer transformer = XmlTransformerFactory.secureXmlTransformer();
         transformer.setOutputProperty("standalone", "yes");
         DOMSource source = new DOMSource(doc);
         transformer.transform(source, outputTarget);
         return outputStream.toByteArray();
      } catch (TransformerException | IOException e) {
         throw new FailureException(e);
      }
   }

   public static void setExtendedAttributeValue(Document doc, String attributeName, String value) {
      NodeList attributes = doc.getElementsByTagName("Attribute");
      if (attributes.getLength() == 0) {
         throw new FailureException("Extended metadata contains no Attribute elements.");
      }

      Node parent = attributes.item(0).getParentNode();
      findAttribute(doc, attributeName).ifPresent(attributex -> attributex.getParentNode().removeChild(attributex));
      Element attribute = doc.createElement("Attribute");
      attribute.setAttribute("name", attributeName);
      Element attributeValue = doc.createElement("Value");
      attributeValue.appendChild(doc.createTextNode(value));
      attribute.appendChild(attributeValue);
      parent.appendChild(doc.createTextNode("\n"));
      parent.appendChild(attribute);
   }

   public static Optional<Node> findAttribute(Document doc, String attributeName) {
      NodeList attributes = doc.getElementsByTagName("Attribute");

      for (int i = 0; i < attributes.getLength(); i++) {
         Node attribute = attributes.item(i);
         if (attributeName.equals(attribute.getAttributes().getNamedItem("name").getNodeValue())) {
            return Optional.of(attribute);
         }
      }

      return Optional.empty();
   }

   private static DocumentBuilderFactory secureDocumentBuilderFactory() throws ParserConfigurationException {
      DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
      factory.setExpandEntityReferences(false);
      factory.setFeature("http://javax.xml.XMLConstants/feature/secure-processing", true);
      factory.setFeature("http://xml.org/sax/features/external-general-entities", false);
      factory.setFeature("http://xml.org/sax/features/external-parameter-entities", false);
      factory.setFeature("http://apache.org/xml/features/nonvalidating/load-external-dtd", false);
      factory.setFeature("http://apache.org/xml/features/nonvalidating/load-dtd-grammar", false);
      return factory;
   }
}