package com.pj.utils;

import org.bouncycastle.asn1.gm.GMNamedCurves;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.CryptoException;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.params.*;
import org.bouncycastle.crypto.signers.SM2Signer;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.jce.interfaces.ECKey;
import org.bouncycastle.jce.interfaces.ECPrivateKey;
import org.bouncycastle.jce.interfaces.ECPublicKey;
import org.bouncycastle.jce.spec.*;
import org.bouncycastle.math.ec.ECCurve;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMWriter;
import org.bouncycastle.openssl.jcajce.JcaPKCS8Generator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedWriter;
import java.io.IOException;
import java.io.StringReader;
import java.io.StringWriter;
import java.math.BigInteger;
import java.nio.charset.Charset;
import java.nio.file.Path;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.nio.file.Files.newBufferedWriter;
import static java.nio.file.Paths.get;
import static java.util.Optional.of;
import static org.bouncycastle.asn1.gm.GMObjectIdentifiers.sm2p256v1;
import static org.bouncycastle.jce.provider.BouncyCastleProvider.PROVIDER_NAME;

/** BC标准实现国密SM2非对称加解密工具类,本工具类的功能包括:
 * <ul>
 * <li>密钥生成
 * <li>密钥恢复
 * </ul>
 * <p>
 * 具体用法，请参考本类对应的 SM2Tests 测试
 * 
 * @author benc.mai
 * @see <a href='https://tools.ietf.org/html/draft-shen-sm2-ecdsa-02'>SM2 Digital Signature Algorithm</a> */
public final class SM2 {

    private static final Logger logger      = LoggerFactory.getLogger(SM2.class);
    public final static String  algorithm   = "EC";
    // private String transformation = Utils.transformation(algorithm, mode, padding);

    public static final String  Pub_File    = SM4.class.getSimpleName() + "PubKey.txt";
    public static final String  Pub_QFile   = SM4.class.getSimpleName() + "PubKeyQ.txt";
    public static final String  Pri_File    = SM4.class.getSimpleName() + "PriKey.txt";
    public static final String  Pri_DFile   = SM4.class.getSimpleName() + "PriKeyD.txt";
    public static final String  Pri_PemFile = SM4.class.getSimpleName() + "PriKey.pem";

    /** 密钥交换字符串编码方式. */
    private Encoder encoder;
    /** 明密文处理字符集. */
    private Charset             charset;

    /** 公钥 JCA/JCE 规范对象形态 */
    private ECPublicKey         pub;
    /** 公钥JCA/JCE 规范字节码形式 */
    private byte[]              pubKeyBytes;
    /** 公钥的 DERX509 编码字符串形式 */
    private String              pubDERX509Str;
    /***/
    private String              pubQStr;

    /** 私钥 JCA/JCE 规范对象形态 */
    private ECPrivateKey        pri;
    /** 私钥JCA/JCE 规范字节码形式 */
    private byte[]              priKeyBytes;
    /** 私钥的 DERPKCS8 编码字符串形式 */
    private String              priDERPKCS8Str;
    /***/
    private String              priDStr;
    /** 私钥的 PEM 编码字符串形式 */
    private String              priPEMStr;

    private SM2(Encoder encoder, Charset charset) {

        this.encoder = of(encoder).get();
        this.charset = of(charset).get();
    }

    //@formatter:off
    public final Encoder getEncoder() { return encoder; }
    public final Charset getCharset() { return charset; }
    public final ECPublicKey getPub() { return pub; }
    public final byte[] getPubKeyBytes() { return pubKeyBytes; }
    public final String getPubDERX509Str() { return pubDERX509Str; }
    public final String getPubQStr() { return pubQStr; }
    public final ECPrivateKey getPri() { return pri; }
    public final byte[] getPriKeyBytes() { return priKeyBytes; }
    public final String getPriDERPKCS8Str() { return priDERPKCS8Str; }
    public final String getPriDStr() { return priDStr; }
    public final String getPriPEMStr() { return priPEMStr; }
    //@formatter:on

    /** 将编码后的密钥对以十六进字符串形式保存到指定目录
     * 
     * @param dir 密钥保存目录,不以 / 结尾
     * @throws RuntimeException */
    public void saveTo(String dir) {

        try {

            Path pDir = get(dir);

            // 公钥
            Path pubPath = pDir.resolve(Pub_File);
            BufferedWriter pubBw = newBufferedWriter(pubPath, charset);
            pubBw.write(pubDERX509Str);
            pubBw.close();
            logger.info("SM2公钥DER X509规范字节码({})保存于:{}", encoder, pubPath);

            Path pubQPath = pDir.resolve(Pub_QFile);
            BufferedWriter pubQBw = newBufferedWriter(pubQPath, charset);
            pubQBw.write(pubQStr);
            pubQBw.close();
            logger.debug("SM2公钥 Q 参数({})保存于:{}", encoder, pubQPath);

            // 私钥
            Path priPath = pDir.resolve(Pri_File);
            BufferedWriter priBw = newBufferedWriter(priPath, charset);
            priBw.write(priDERPKCS8Str);
            priBw.close();
            logger.debug("SM2私钥DER PKCS8规范字节码({})保存于:{}", encoder, priPath);

            Path priDPath = pDir.resolve(Pri_DFile);
            BufferedWriter priDBw = newBufferedWriter(priDPath, charset);
            priDBw.write(priDStr);
            priDBw.close();
            logger.debug("SM2私钥 D 参数({})保存于:{}", encoder, priDPath);

            Path priPemPath = pDir.resolve(Pri_PemFile);
            BufferedWriter PriPemBw = newBufferedWriter(priPemPath, charset);
            PriPemBw.write(priPEMStr);
            PriPemBw.close();
            logger.debug("SM2私钥 PEM 规范编码保存于:{}", priPemPath);
        } catch (Exception e) {
            logger.warn("SM2 密钥保存失败.原因: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

    @Override
    public String toString() {

        return super.toString();
        // return String.format("SM2封装对象 - Transformation: %s Key(%s %s byte):%s Charset:%s", transformation, encoder,
        // keyBytes.length, encodedKeyStr, charset);
    }

    // 静态工具方法
    // -----------------------------------------------------------------------------------------------------------------

    public static SM2 genKeyPair() {

        return genKeyPair(Encoder.Hex, UTF_8);
    }

    /** 生成SM2密钥对.密钥长度 应为 256
     * 
     * @param encoder 密钥、明/密文字节码编码器
     * @param charset 字符集
     * @return SM2 封装对象
     * @throws RuntimeException */
    public static SM2 genKeyPair(Encoder encoder, Charset charset) {

        SM2 k = new SM2(encoder, charset);

        try {
            // 产生密钥对
            KeyPairGenerator kpg = KeyPairGenerator.getInstance(algorithm, PROVIDER_NAME);
            kpg.initialize(new ECNamedCurveGenParameterSpec(sm2p256v1.getId()));// ("sm2p256v1"));
            KeyPair keyPair = kpg.generateKeyPair();

            // 公钥
            k.pub = (BCECPublicKey) keyPair.getPublic();
            k.pubKeyBytes = encodePubKeyWithX509(k.pub);
            k.pubDERX509Str = k.encoder.encode(k.pubKeyBytes);
            k.pubQStr = k.encoder.encode(k.pub.getQ().getEncoded(Boolean.FALSE));

            // 私钥
            k.pri = (BCECPrivateKey) keyPair.getPrivate();
            k.priKeyBytes = encodePriKeyWithPKCS8(k.pri);
            k.priDERPKCS8Str = k.encoder.encode(k.priKeyBytes);
            k.priDStr = k.encoder.encode(k.pri.getD().toByteArray());
            k.priPEMStr = encodePriKey2PEM(k.pri);

        } catch (Exception e) {
            logger.warn("SM2 密钥对生成失败.原因: {}", e.getMessage());
            throw new RuntimeException(e);
        }

        return k;
    }

    // DER编码密钥对
    //
    // JCE本身是支持DER编码密钥对的解析的，可以参见PKCS8EncodedKeySpec和X509EncodedKeySpec
    // DER编码是ASN.1编码规则中的一个子集，具体格式如何编排没有去了解，但最终呈现肯定的是一堆有规律的二进制组合而成
    // PKCS#8定义了私钥信息语法和加密私钥语法，而X509定义证书规范，通常都会用DER和PEM进行编码存储，而在JAVA中则使用的DER
    // ------------------------------------------------------------------------------------------------------------------
    /** 生成公匙X509编码字节数组 */
    public static byte[] encodePubKeyWithX509(PublicKey key) {

        // key.getFormat();
        // PublicKey pubKey = /* ( BCRSAPublicKey | RSAPublicKey) */key;
        byte[] pubKeyBytes = key.getEncoded();// X509EncodedKeySpec
        // RSAPublicKeySpec pubKeySpec = new RSAPublicKeySpec(pubKey.getModulus(), pubKey.getPublicExponent());
        return pubKeyBytes;
    }

    /** 生成私匙PKCS#8编码字节数组
     * 
     * @throws IOException */
    public static byte[] encodePriKeyWithPKCS8(PrivateKey key) {

        // key.getFormat();
        // PrivateKey priKey = /* ( BCRSAPrivateCrtKey | RSAPrivateKey) */key;
        byte[] priKeyBytes = key.getEncoded();// PKCS8EncodedKeySpec
        // logger.debug("RSAPrivateKey:{}", ByteUtilities.asHex(priKeyBytes));
        return priKeyBytes;
    }

    /** 从X.509编码字节数组恢复公钥
     * 
     * @param x590Key DER编码公钥字节码
     * @return 公钥
     * @throws RuntimeException */
    public static ECPublicKey recoveryPubKeyFromX509EncodedKeySpec(byte[] x590Key) {

        try {
            KeySpec keySpec = new X509EncodedKeySpec(x590Key);
            KeyFactory keyFactory = KeyFactory.getInstance(algorithm);
            ECPublicKey pubKey = (ECPublicKey) keyFactory.generatePublic(keySpec);
            return pubKey;
        } catch (GeneralSecurityException e) {
            logger.warn("从X.509编码字节数组恢复  SM2 公钥失败.原因: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

    public static ECPublicKey recoveryPubKeyFromQ(byte[] pubQbytes) {

        // ECDomainParameters d_params = new ECDomainParameters(curve, G, BigInteger.ONE);
        // ECPublicKeyParameters spec = new ECPublicKeyParameters(Q, d_params);

        X9ECParameters x9 = GMNamedCurves.getByOID(sm2p256v1);
        ECCurve curve = x9.getCurve();
        ECPoint G = x9.getG();
        BigInteger N = x9.getN();

        ECPublicKey pubKey;
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(algorithm, PROVIDER_NAME);
            ECParameterSpec paramSpec = new ECNamedCurveParameterSpec(sm2p256v1.getId(), curve, G, N);
            // byte[] Qbytes = encoder.decode(pub_Q_Str);
            ECPoint Q = curve.decodePoint(pubQbytes);
            KeySpec keySpec = new ECPublicKeySpec(Q, paramSpec);
            pubKey = (ECPublicKey) keyFactory.generatePublic(keySpec);
            return pubKey;
        } catch (NoSuchAlgorithmException | NoSuchProviderException | InvalidKeySpecException e) {
            logger.warn("从 Q 参数恢复 SM2 公钥失败.原因: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

    /** 从PKCS#8字节数组恢复私钥
     * 
     * @param pkcs8Key DER编码私钥字节码
     * @return 私钥
     * @throws RuntimeException */
    public static ECPrivateKey recoveryPriKeyFromPKCS8EncodedKeySpec(byte[] pkcs8Key) {

        try {
            KeySpec keySpec = new PKCS8EncodedKeySpec(pkcs8Key);
            KeyFactory keyFactory = KeyFactory.getInstance(algorithm);
            ECPrivateKey privKey = (ECPrivateKey) keyFactory.generatePrivate(keySpec);
            return privKey;
        } catch (GeneralSecurityException e) {
            logger.warn("从PKCS#8字节数组恢复 SM2 私钥失败.原因: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

    public static ECPrivateKey recoveryPriKeyFromD(byte[] Dbytes) {

        X9ECParameters x9 = GMNamedCurves.getByOID(sm2p256v1);
        ECCurve curve = x9.getCurve();
        ECPoint G = x9.getG();
        BigInteger N = x9.getN();

        ECPrivateKey priKey;
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(algorithm, PROVIDER_NAME);
            ECParameterSpec paramSpec = new ECNamedCurveParameterSpec(sm2p256v1.getId(), curve, G, N);

            BigInteger D = new BigInteger(1, Dbytes);
            KeySpec keySpec = new ECPrivateKeySpec(D, paramSpec);
            priKey = (ECPrivateKey) keyFactory.generatePrivate(keySpec);
            return priKey;
        } catch (NoSuchAlgorithmException | NoSuchProviderException | InvalidKeySpecException e) {
            logger.warn("从 D 参数恢复 SM2 私钥失败.原因: {}", e.getMessage());
            throw new RuntimeException(e);
        }

    }

    // PEM编码密钥恢复
    //
    //
    // ------------------------------------------------------------------------------------------------------------------
    /** 将私钥转换成PEM字符串
     * 
     * @param key
     * @return 私钥
     * @throws RuntimeException */
    public static String encodePriKey2PEM(PrivateKey key) {

        try {
            StringWriter sw = new StringWriter();
            JcaPEMWriter pw = new JcaPEMWriter(sw);
            pw.writeObject(new JcaPKCS8Generator(key, null));
            // pw.writeObject(key);
            pw.close();
            String priKeyPEM = sw.toString();
            return priKeyPEM;
        } catch (IOException e) {
            logger.warn("将 SM2 私钥编码成 PEM 格式字符串失败.原因: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

    /** 从PEM 字符串恢复出私钥
     * 
     * @param PEM
     * @return 私钥
     * @throws RuntimeException */
    public static ECPrivateKey recoveryPriKeyFromPEM(String PEM) {

        try {
            PEMParser pr = new PEMParser(new StringReader(PEM));
            PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(pr.readPemObject().getContent());
            final KeyFactory keyFactory = KeyFactory.getInstance(algorithm);
            ECPrivateKey privKey = (ECPrivateKey) keyFactory.generatePrivate(pkcs8KeySpec);
            pr.close();
            return privKey;
        } catch (Exception e) {
            logger.warn("从 PEM 格式字符串恢复 SM2 私钥失败.原因: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

    public static byte[] encrypt(ECKey ecKey, byte[] plainBytes) {

        ECKeyParameters key_param = getECKeyParameters(ecKey);

        SM2Engine engine = new SM2Engine();
        engine.init(true, new ParametersWithRandom(key_param, new SecureRandom()));
        byte[] cipherBtyes;
        try {
            cipherBtyes = engine.processBlock(plainBytes, 0, plainBytes.length);
            return cipherBtyes;
        } catch (InvalidCipherTextException e) {
            logger.warn("SM2 {} 加密失败.原因: {}", ecKey.getClass().getSimpleName(), e.getMessage());
            throw new RuntimeException(e);
        }

    }

    public static byte[] decrypt(ECKey ecKey, byte[] cipherBtyes) {

        ECKeyParameters key_param = getECKeyParameters(ecKey);

        SM2Engine engine = new SM2Engine();
        engine.init(false, key_param);
        byte[] plainBytes;
        try {
            plainBytes = engine.processBlock(cipherBtyes, 0, cipherBtyes.length);
            return plainBytes;
        } catch (InvalidCipherTextException e) {
            logger.warn("SM2 {} 解密失败.原因: {}", ecKey.getClass().getSimpleName(), e.getMessage());
            throw new RuntimeException(e);
        }

    }

    /** ECC私钥签名
     *
     * @param ecKey ECC私钥
     * @param withId 可以为null，若为null，则默认withId为字节数组:"1234567812345678".getBytes()
     * @param srcData 源数据
     * @return 签名
     * @throws NoSuchAlgorithmException
     * @throws NoSuchProviderException
     * @throws CryptoException */
    public static byte[] sign(ECKey ecKey, byte[] withId, byte[] srcData) {

        ECKeyParameters key_param = getECKeyParameters(ecKey);

        SM2Signer signer = new SM2Signer();
        CipherParameters param = null;
        ParametersWithRandom pwr = new ParametersWithRandom(key_param, new SecureRandom());
        if (withId != null) {
            param = new ParametersWithID(pwr, withId);
        } else {
            param = pwr;
        }
        signer.init(true, param);
        signer.update(srcData, 0, srcData.length);
        try {
            return signer.generateSignature();
        } catch (CryptoException e) {
            logger.warn("SM2 {} 签名失败.原因: {}", ecKey.getClass().getSimpleName(), e.getMessage());
            throw new RuntimeException(e);
        }
    }

    /** ECC公钥验签
     *
     * @param ecKey ECC公钥
     * @param withId 可以为null，若为null，则默认withId为字节数组:"1234567812345678".getBytes()
     * @param srcData 源数据
     * @param sign 签名
     * @return 验签成功返回true，失败返回false */
    public static boolean verify(ECKey ecKey, byte[] withId, byte[] srcData, byte[] sign) {

        ECKeyParameters key_param = getECKeyParameters(ecKey);

        SM2Signer signer = new SM2Signer();
        CipherParameters param = null;
        if (withId != null) {
            param = new ParametersWithID(key_param, withId);
        } else {
            param = key_param;
        }
        signer.init(false, param);
        signer.update(srcData, 0, srcData.length);
        return signer.verifySignature(sign);
    }

    private static ECKeyParameters getECKeyParameters(ECKey ecKey) {

        ECParameterSpec param_spec = ecKey.getParameters();
        ECDomainParameters domain_param = new ECDomainParameters(param_spec.getCurve(), param_spec.getG(),
                param_spec.getN());

        ECKeyParameters key_param = null;
        if (ecKey instanceof ECPublicKey)
            key_param = new ECPublicKeyParameters(((ECPublicKey) ecKey).getQ(), domain_param);
        else if (ecKey instanceof ECPrivateKey)
            key_param = new ECPrivateKeyParameters(((ECPrivateKey) ecKey).getD(), domain_param);

        return key_param;
    }
}
