XMLSignatureUtil v2

import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import javax.xml.XMLConstants;
import javax.xml.crypto.AlgorithmMethod;
import javax.xml.crypto.KeySelector;
import javax.xml.crypto.KeySelectorException;
import javax.xml.crypto.KeySelectorResult;
import javax.xml.crypto.XMLCryptoContext;
import javax.xml.crypto.XMLStructure;
import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.Transform;
import javax.xml.crypto.dsig.XMLSignature;
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 javax.xml.namespace.NamespaceContext;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Source;
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 javax.xml.transform.stream.StreamSource;
import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathConstants;
import javax.xml.xpath.XPathExpression;
import javax.xml.xpath.XPathExpressionException;
import javax.xml.xpath.XPathFactory;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.StringWriter;
import java.nio.charset.StandardCharsets;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;

public class XMLSignatureUtil {

    private static final DocumentBuilderFactory DOCUMENT_BUILDER_FACTORY = createDocumentBuilderFactory();
    private static final KeySelector KEY_SELECTOR = new KeyValueKeySelector();
    private static final XMLSignatureFactory XML_SIGNATURE_FACTORY = XMLSignatureFactory.getInstance("DOM", new org.apache.jcp.xml.dsig.internal.dom.XMLDSigRI());

    public static Document createDocument(InputStream inputStream) throws Exception {
        try (inputStream) {
            DocumentBuilder documentBuilder = createDocumentBuilder();
            return documentBuilder.parse(inputStream);
        }
    }

    public static boolean isValid(String xml) throws Exception {
        try (InputStream sanitizedStream = new ByteArrayInputStream(xml.getBytes(StandardCharsets.UTF_8))) {
            return isValid(sanitizedStream);
        }
    }

    public static boolean isValid(InputStream inputStream) throws Exception {
        try (InputStream sanitizedStream = sanitizeAndConvertToStream(inputStream)) {
            Document document = createDocument(sanitizedStream);

            Element elementWithId = findElementWithId(document);
            elementWithId.setIdAttribute("Id", true);

            NodeList signatureElement = document.getElementsByTagNameNS(XMLSignature.XMLNS, "Signature");
            if (signatureElement.getLength() == 0) {
                throw new Exception("Elemento de assinatura não encontrado.");
            }

            NodeList keyInfoList = document.getElementsByTagNameNS(XMLSignature.XMLNS, "KeyInfo");
            if (keyInfoList.getLength() == 0) {
                throw new Exception("Elemento KeyInfo não encontrado.");
            }

            DOMValidateContext domValidateContext = new DOMValidateContext(KEY_SELECTOR, signatureElement.item(0));
            XMLSignature signature = XML_SIGNATURE_FACTORY.unmarshalXMLSignature(domValidateContext);

            return signature.validate(domValidateContext);
        }
    }

    public static String signed(String xml, String password, byte[] certificate) throws Exception {
        try (InputStream sanitizedStream = new ByteArrayInputStream(xml.getBytes(StandardCharsets.UTF_8))) {
            return signed(sanitizedStream, password, certificate);
        }
    }

    public static String signed(InputStream inputStream, String password, byte[] pfxPath) throws Exception {
        Document document = createDocument(inputStream);
        KeyStore keyStore = KeyStore.getInstance("PKCS12");
        try (var keyStream = new ByteArrayInputStream(pfxPath)) {
            keyStore.load(keyStream, password.toCharArray());
        }

        String alias = keyStore.aliases().nextElement();
        PrivateKey privateKey = (PrivateKey) keyStore.getKey(alias, password.toCharArray());
        X509Certificate certificate = (X509Certificate) keyStore.getCertificate(alias);

        Element elementToSign = findElementWithId(document);
        elementToSign.setIdAttribute("Id", true);
        String id = elementToSign.getAttribute("Id");

        // Referência com as transformações exigidas
        Reference ref = XML_SIGNATURE_FACTORY.newReference(
                "#" + id,
                XML_SIGNATURE_FACTORY.newDigestMethod(DigestMethod.SHA256, null),
                Arrays.asList(
                        XML_SIGNATURE_FACTORY.newTransform(Transform.ENVELOPED, (TransformParameterSpec) null),
                        XML_SIGNATURE_FACTORY.newTransform(CanonicalizationMethod.INCLUSIVE, (C14NMethodParameterSpec) null)
                ),
                null,
                null
        );

        // Informação assinada
        SignedInfo signedInfo = XML_SIGNATURE_FACTORY.newSignedInfo(
                XML_SIGNATURE_FACTORY.newCanonicalizationMethod(CanonicalizationMethod.INCLUSIVE, (C14NMethodParameterSpec) null),
                XML_SIGNATURE_FACTORY.newSignatureMethod(SignatureMethod.RSA_SHA256, null),
                Collections.singletonList(ref)
        );

        // Adicionar informações da chave
        KeyInfoFactory keyInfoFactory = XML_SIGNATURE_FACTORY.getKeyInfoFactory();
        X509Data x509Data = keyInfoFactory.newX509Data(Collections.singletonList(certificate));
        KeyInfo keyInfo = keyInfoFactory.newKeyInfo(Collections.singletonList(x509Data));

        DOMSignContext domSignContext = new DOMSignContext(privateKey, document.getDocumentElement());

        XMLSignature xmlSignature = XML_SIGNATURE_FACTORY.newXMLSignature(signedInfo, keyInfo);
        xmlSignature.sign(domSignContext);

        return transformDocumentToString(document);
    }

    private static String transformDocumentToString(Document document) throws TransformerException {
        TransformerFactory tf = TransformerFactory.newInstance();
        Transformer transformer = tf.newTransformer();
        StringWriter writer = new StringWriter();
        transformer.transform(new DOMSource(document), new StreamResult(writer));
        return writer.toString();
    }

    private static synchronized DocumentBuilder createDocumentBuilder() throws ParserConfigurationException {
        return DOCUMENT_BUILDER_FACTORY.newDocumentBuilder();
    }

    private static DocumentBuilderFactory createDocumentBuilderFactory() {
        DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
        factory.setNamespaceAware(true);
        try {
            factory.setFeature(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/disallow-doctype-decl", true);
        } catch (ParserConfigurationException e) {
            throw new IllegalStateException("Erro ao configurar DocumentBuilderFactory", e);
        }
        return factory;
    }

    private static InputStream sanitizeAndConvertToStream(InputStream inputStream) throws TransformerException, IOException {
        try (StringWriter writer = new StringWriter()) {
            TransformerFactory transformerFactory = TransformerFactory.newInstance();
            transformerFactory.setFeature(XMLConstants.FEATURE_SECURE_PROCESSING, true);
            Transformer transformer = transformerFactory.newTransformer();
            transformer.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "yes");
            transformer.setOutputProperty(OutputKeys.INDENT, "no");
            transformer.setOutputProperty(OutputKeys.METHOD, "xml");
            Source source = new StreamSource(inputStream);
            transformer.transform(source, new StreamResult(writer));
            String sanitizedXml = writer.toString()
                    .replaceAll("\\s+", " ")
                    .replaceAll("\\s*>\\s*", ">")
                    .replaceAll("\\s*<\\s*", "<")
                    .trim();
            return new ByteArrayInputStream(sanitizedXml.getBytes(StandardCharsets.UTF_8));
        }
    }

    private static Element findFirstElementChild(Node parent) {
        if (parent == null) return null;
        NodeList children = parent.getChildNodes();
        for (int i = 0; i < children.getLength(); i++) {
            Node child = children.item(i);
            if (child.getNodeType() == Node.ELEMENT_NODE) {
                return (Element) child;
            }
        }
        return null;
    }

    private static Element findElementWithId(Document document) throws Exception {
        XPathFactory xpathFactory = XPathFactory.newInstance();
        XPath xpath = xpathFactory.newXPath();

        xpath.setNamespaceContext(new NamespaceContext() {
            @Override
            public String getNamespaceURI(String prefix) {
                return switch (prefix.toLowerCase()) {
                    case "nfe" -> "http://www.portalfiscal.inf.br/nfe";
                    case "cte" -> "http://www.portalfiscal.inf.br/cte";
                    default -> XMLConstants.NULL_NS_URI;
                };
            }

            @Override
            public String getPrefix(String namespaceURI) {
                return null;
            }

            @Override
            public Iterator<String> getPrefixes(String namespaceURI) {
                return Collections.emptyIterator();
            }
        });

        try {
            XPathExpression expr = xpath.compile("//*[starts-with(@Id, 'NFe') or starts-with(@Id, 'ID') or starts-with(@Id, 'CTe')]");
            NodeList nodes = (NodeList) expr.evaluate(document, XPathConstants.NODESET);
            if (nodes.getLength() == 0) {
                throw new Exception("Elemento com Id não encontrado. Estrutura do XML:\n");
            }
            return (Element) nodes.item(0);
        } catch (XPathExpressionException e) {
            throw new Exception("Erro ao buscar elemento com Id", e);
        }
    }

    // Classe interna permanece inalterada, exceto para o uso de exceções específicas.
    private static class KeyValueKeySelector extends KeySelector {

        public KeySelectorResult select(KeyInfo keyInfo, KeySelector.Purpose purpose, AlgorithmMethod method, XMLCryptoContext context) throws KeySelectorException {
            if (keyInfo == null) {
                throw new KeySelectorException("Nenhum KeyInfo encontrado!");
            }
            for (XMLStructure info : keyInfo.getContent()) {
                if (info instanceof X509Data) {
                    try {
                        for (Object data : ((X509Data) info).getContent()) {
                            if (data instanceof X509Certificate) {
                                PublicKey pk = ((X509Certificate) data).getPublicKey();
                                return new SimpleKeySelectorResult(pk);
                            }
                        }
                    } catch (Exception e) {
                        throw new KeySelectorException(e);
                    }
                }
            }
            throw new KeySelectorException("Nenhum KeyValue encontrado!");
        }
    }

    private record SimpleKeySelectorResult(PublicKey pk) implements KeySelectorResult {
        public PublicKey getKey() {
            return pk;
        }
    }
}