package com.zbkj.service.wangshang.api.internal.util;

import com.zbkj.service.wangshang.api.MybankApiException;
import com.zbkj.service.wangshang.api.MybankApiExceptionEnum;
import com.zbkj.service.wangshang.api.MybankConstants;
import org.apache.commons.io.IOUtils;
import org.apache.xml.security.exceptions.Base64DecodingException;
import org.apache.xml.security.exceptions.XMLSecurityException;
import org.apache.xml.security.signature.XMLSignature;
import org.apache.xml.security.transforms.Transforms;
import org.apache.xml.security.utils.Base64;
import org.apache.xml.security.utils.Constants;
import org.apache.xml.security.utils.XMLUtils;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import java.io.ByteArrayOutputStream;
import java.io.UnsupportedEncodingException;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * XML签名/验签工具类，包含以下功能
 * <pre>
 * 1. 获取RSA的公钥
 * 2. 获取RSA的私钥
 * 3. 进行RSA签名
 * 4. 进行RSA验签
 * </pre>
 */
public class SignatureUtils {
    static {
        org.apache.xml.security.Init.init();
    }

    /**
     * 根据String获取私钥 - RSA
     *
     * @param encodePrivateKey
     * @return
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public static PrivateKey getPrivateKey(String encodePrivateKey, String signType)
            throws NoSuchAlgorithmException,
            InvalidKeySpecException, Base64DecodingException {
        byte[] privateKeyBytes = Base64.decode(encodePrivateKey);
        KeyFactory kf = KeyFactory.getInstance(signType);
        PrivateKey privateKey = kf.generatePrivate(new PKCS8EncodedKeySpec(privateKeyBytes));
        return privateKey;
    }

    /**
     * 根据String获取公钥 - RSA
     *
     * @param encodePublicKey
     * @return
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public static PublicKey getPublicKey(String encodePublicKey, String signType) throws NoSuchAlgorithmException,
            InvalidKeySpecException, Base64DecodingException {
        byte[] publicKeyKeyBytes = Base64.decode(encodePublicKey);
        KeyFactory kf = KeyFactory.getInstance(signType);
        PublicKey publicKey = kf.generatePublic(new X509EncodedKeySpec(publicKeyKeyBytes));
        return publicKey;
    }

    /**
     * XML签名
     *
     * @param privateKey          私钥数据，PKCS#8编码格式
     * @param xmlDocument         XML文件内容， UTF-8编码
     * @param elementTagName      续签签名的Tag名称
     * @param algorithm           签名算法 {@link XMLSignature} 支持下列算法
     *                            <ul>
     *                            <li>XMLSignature.ALGO_ID_SIGNATURE_RSA</li>
     *                            <li>XMLSignature.ALGO_ID_SIGNATURE_RSA_SHA1</li>
     *                            <li>XMLSignature.ALGO_ID_SIGNATURE_RSA_SHA256</li>
     *                            <li>XMLSignature.ALGO_ID_SIGNATURE_RSA_SHA384</li>
     *                            <li>XMLSignature.ALGO_ID_SIGNATURE_RSA_SHA512</li>
     *                            </ul>
     * @param signatureAppendMode 签名节点的附加模式
     *                            <ul>
     *                            <li>作为子节点： XmlSignatureAppendMode.AS_CHILDREN</li>
     *                            <li>作为兄弟节点：XmlSignatureAppendMode.AS_BROTHER</li>
     *                            </ul>
     * @return 签名后的XML string
     * @throws MybankApiException, XMLSecurityException
     */
    public static String signXmlElement(PrivateKey privateKey, Document xmlDocument,
                                        String elementTagName, String algorithm,
                                        int signatureAppendMode) throws MybankApiException, XMLSecurityException {
        XMLSignature xmlSignature = new XMLSignature(xmlDocument, xmlDocument.getDocumentURI(),
                algorithm);

        NodeList nodeList = xmlDocument.getElementsByTagName(elementTagName);
        if (nodeList == null || nodeList.getLength() - 1 < 0) {
            throw new MybankApiException(MybankApiExceptionEnum.DOCUMENT_ELEMENT_NOT_EXIST);
        }

        Node elementNode = nodeList.item(0);
        if (elementNode == null) {
            throw new MybankApiException(MybankApiExceptionEnum.DOCUMENT_ELEMENT_NOT_EXIST);
        }

        elementNode.appendChild(xmlSignature.getElement());
        if (signatureAppendMode == XmlSignatureAppendMode.AS_CHILDREN) {
            elementNode.appendChild(xmlSignature.getElement());
        } else if (signatureAppendMode == XmlSignatureAppendMode.AS_BROTHER) {
            elementNode.getParentNode().appendChild(xmlSignature.getElement());
        } else {
            throw new MybankApiException(MybankApiExceptionEnum.ILLEGAL_APPEND_MODE);
        }

        Transforms transforms = new Transforms(xmlDocument);
        transforms.addTransform(Transforms.TRANSFORM_ENVELOPED_SIGNATURE);
        xmlSignature.addDocument(MybankConstants.NULL_STRING, transforms, Constants.ALGO_ID_DIGEST_SHA1);

        xmlSignature.sign(privateKey);

        ByteArrayOutputStream os = null;
        try {
            os = new ByteArrayOutputStream();
            XMLUtils.outputDOM(xmlDocument, os);
            return os.toString(MybankConstants.CHARSET_UTF8);
        } catch (UnsupportedEncodingException e) {
            throw new MybankApiException(e);
        } finally {
            IOUtils.closeQuietly(os);
        }
    }

    /**
     * 验证XML签名
     *
     * @param publicKey   公钥数据 X509编码
     * @param xmlDocument XML内容
     * @return 签名验证结果 boolean
     * @throws MybankApiException, XMLSecurityException
     */
    public static boolean verifyXmlElement(PublicKey publicKey, Document xmlDocument)
            throws MybankApiException, XMLSecurityException {
        NodeList signatureNodes = xmlDocument.getElementsByTagNameNS(Constants.SignatureSpecNS,
                MybankConstants.SIGNATURE);
        if (signatureNodes == null || signatureNodes.getLength() < 1) {
            throw new MybankApiException(MybankApiExceptionEnum.SIGNATURE_ELEMENT_NOT_EXIST);
        }

        Element signElement = (Element) signatureNodes.item(0);
        if (signElement == null) {
            throw new MybankApiException(MybankApiExceptionEnum.SIGNATURE_ELEMENT_NOT_EXIST);
        }

        XMLSignature signature = new XMLSignature(signElement, MybankConstants.NULL_STRING);
        return signature.checkSignatureValue(publicKey);
    }
}