package com.pj.utils;

import cn.hutool.crypto.Mode;
import cn.hutool.crypto.Padding;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.BufferedWriter;
import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.file.Path;
import java.security.GeneralSecurityException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.SecureRandom;

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.jce.provider.BouncyCastleProvider.PROVIDER_NAME;

/** BC标准实现国密SM4对称加解密工具类,本工具类的功能包括:
 * <ul>
 * <li>密钥生成
 * <li>密钥恢复
 * <li>加密.密文自动串联 iv长度、iv字节码 、密文字节码
 * <li>解密.密文自动提取 iv长度、iv字节码 、密文字节码
 * <li>支持全部SM4支持的区块工作模式和填充模式及扩展模式
 * </ul>
 * <p>
 * 具体用法，请参考本类对应的 SM4Tests 测试
 * 
 * @author benc.mai
 * @see <a href='https://tools.ietf.org/id/draft-ribose-cfrg-sm4-01.html'>The SM4 Block Cipher Algorithm And Its Modes
 *      Of Operations</a> */
public final class SM4 {

    private static final Logger logger         = LoggerFactory.getLogger(SM4.class);
    public static final String  Key_File_Name  = SM4.class.getSimpleName() + "Key.txt";
    public static final String  algorithm      = "SM4";

    /** 算法名 */

    /** 使用何种区块加解密工作模式 .默认： CBC */
    private Mode mode           = Mode.CBC;
    /** 使用何种区块加解密填充方式.默认： PKCS5Padding */
    private Padding padding        = Padding.PKCS5Padding;

    /** 算法变种,用于 Cipher 加解密过程处理.
     * <p>
     * 格式为:
     * 算法/模式/填充 (algorithm/mode/padding)
     * <p>
     * SM4 工作模式包含 ECB、CBC、CFBx、OFBx、CTR.(ECB、CBC 较为常用,除无工作模式和ECB外,其他模式都必须引入IV初始向量或额外变量)
     * <p>
     * SM4 在不同工作模式下支持的算法变种 :
     * <li>SM4
     * <li>SM4/ECB/PKCS5Padding、PKCS7Padding
     * <li>SM4/CBC/PKCS5Padding、PKCS7Padding、ISO10126Padding
     * <li>SM4/CFBx/NoPadding、PKCS5Padding、PKCS7Padding、ISO10126Padding
     * <li>SM4/OFBx/NoPadding、PKCS5Padding、PKCS7Padding、ISO10126Padding
     * <li>SM4/CTR/NoPadding、PKCS5Padding、PKCS7Padding、ISO10126Padding
     * <li>SM4/EAX/NoPadding
     * <li>SM4/GCM/NoPadding */
    private String              transformation = Utils.transformation(algorithm, mode, padding);

    /** 密钥交换字符串编码方式.默认：Hex */
    private Encoder encoder;
    /** 明密文处理字符集.默认： UTF-8 */
    private Charset             charset;

    /** 密钥JCA/JCE规范对象形态 */
    private SecretKey           key;
    /** 密钥JCA/JCE规范字节码形态 */
    private byte[]              keyBytes;
    /** 密钥编码字符串形态 */
    private String              encodedKeyStr;

    /** 构造SM4封装对象
     * 
     * @param encoder 密钥、明/密文字节码编码器
     * @param charset 字符集
     * @param mode 区块加解密工作模式
     * @param padding 区块加解密填充方式 */
    private SM4(Encoder encoder, Charset charset, Mode mode, Padding padding) {

        this.encoder = of(encoder).get();
        this.charset = of(charset).get();
        this.mode = of(mode).get();
        this.padding = of(padding).get();
        this.transformation = Utils.transformation(algorithm, mode, padding);
    }

    //@formatter:off
    public Encoder getEncoder() { return encoder; }
    public Charset getCharset() { return charset; }
    public Mode getMode() { return mode; }
    public Padding getPadding() { return padding; }
    public String getTransformation() { return transformation;}
    public SecretKey getKey() { return key; }
    public byte[] getKeyBytes() { return keyBytes; }
    public String getEncodedKeyStr() { return encodedKeyStr; }
    //@formatter:on

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

        Path pDir = get(dir);

        try {
            BufferedWriter bw = newBufferedWriter(pDir.resolve(Key_File_Name), UTF_8);
            bw.write(encodedKeyStr);
            bw.close();
            logger.info("保存 SM4密钥文件于:{}", pDir.resolve(Key_File_Name));

        } catch (IOException e) {
            logger.warn("SM4 密钥保存失败.原因: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

    @Override
    public String toString() {

        return String.format("SM4 - Transformation: %s Key(%s %s byte):%s Charset:%s", transformation, encoder,
                keyBytes.length, encodedKeyStr, charset);
    }

    // 静态工具方法
    // -----------------------------------------------------------------------------------------------------------------
    public static SM4 genKey() {

        return genKey(Encoder.Hex, UTF_8, Mode.CBC, Padding.PKCS5Padding);
    }

    /** 生成SM4封装对象
     * 
     * @param encoder 密钥、明/密文字节码编码器
     * @param charset 字符集
     * @param mode 区块加解密工作模式
     * @param padding 区块加解密填充方式
     * @return SM4 封装对象
     * @throws RuntimeException */
    public static SM4 genKey(Encoder encoder, Charset charset, Mode mode, Padding padding) {

        SM4 k = new SM4(encoder, charset, mode, padding);

        try {
            KeyGenerator kg = KeyGenerator.getInstance(algorithm, PROVIDER_NAME);
            SecureRandom rand = new SecureRandom();
            // keyGenerator.init(128, rand);// SM4 标准实现密钥长度必须 128 bit
            kg.init(rand);

            k.key = kg.generateKey();
            k.keyBytes = k.key.getEncoded();
            k.encodedKeyStr = k.encoder.encode(k.keyBytes);
            logger.info(k.toString());
            return k;
        } catch (NoSuchAlgorithmException | NoSuchProviderException e) {
            logger.warn("SM4 Key Generation Fail.Reason: {}", e.getMessage());
            throw new RuntimeException(e);
        }

    }

    public static SM4 recovery(String encodedKeyStr) {

        return recovery(Encoder.Hex, UTF_8, encodedKeyStr, Mode.CBC, Padding.PKCS5Padding);
    }

    /** 还原密钥封装对象
     * 
     * @param encoder 密钥、明/密文字节码编码器
     * @param charset 字符集
     * @param encodedKeyStr JCA/JCE 规范字节码编码后的密钥字符串
     * @param mode 区块加解密工作模式
     * @param padding 区块加解密填充方式
     * @return SM4 封装对象 */
    public static SM4 recovery(Encoder encoder, Charset charset, String encodedKeyStr, Mode mode, Padding padding) {

        SM4 k = new SM4(encoder, charset, mode, padding);

        k.key = new SecretKeySpec(encoder.decode(encodedKeyStr), k.transformation);
        k.keyBytes = k.key.getEncoded();
        k.encodedKeyStr = encodedKeyStr;
        logger.info(k.toString());
        return k;
    }

    /** 只还原 JCA/JCE 密钥对象
     * 
     * @param keyBytes JCA/JCE 规范密钥字节码 */
    public static SecretKey recovery(byte[] keyBytes) {

        SecretKey key = new SecretKeySpec(keyBytes, algorithm);
        return key;
    }

    /** SM4 加密
     * 
     * @param keyBytes 密钥字节码
     * @param plainBytes 待加密原文字节码
     * @param transformation 算法变种
     * @return 密文字节码
     * @throws RuntimeException */
    public static byte[] encrypt(byte[] keyBytes, byte[] plainBytes, String transformation) {

        try {
            Cipher cipher = Cipher.getInstance(of(transformation).get(), PROVIDER_NAME);
            cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(of(keyBytes).get(), transformation));

            byte[] cipherBytes = cipher.doFinal(of(plainBytes).get());

            // 特定的区块加密工作模式,除ECB外其他的一般都会产生初始向量(IV).IV的长度一般和区块长度一致
            // 为便于传输，可以把IV作为前缀或后缀串联到密文字节码上，这里默认串联到密文数组前
            // 解密时必须先分别提取IV和密文数组字节码
            // SM4算法不同的区块工作模式IV长度由12-15-16字节不等,大部分都为16字节
            // IV长度(1bit)+ IV字节码(12-16bit)+密文
            byte[] ivBytes = cipher.getIV();
            if (ivBytes != null) {
                // IvParameterSpec iv_spec = new IvParameterSpec(iv);
                // Mode.SHORT_IV_Okay.contains(mode);
                // logger.info("加密方初始向量(IV {} bytes): {}", ivBytes.length, Encoder.Hex.encode(ivBytes));
                byte[] combineIvlenghtAndIvAndCipher = new byte[1 + ivBytes.length + cipherBytes.length];
                combineIvlenghtAndIvAndCipher[0] = (byte) ivBytes.length;
                System.arraycopy(ivBytes, 0, combineIvlenghtAndIvAndCipher, 1, ivBytes.length);
                System.arraycopy(cipherBytes, 0, combineIvlenghtAndIvAndCipher, ivBytes.length + 1, cipherBytes.length);
                cipherBytes = combineIvlenghtAndIvAndCipher;
            }
            return cipherBytes;
        } catch (GeneralSecurityException e) {
            logger.warn("SM4 Encrypt.Reason: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

    /** SM4 解密
     * 
     * @param keyBytes 密钥字节码
     * @param cipherBytes 密文字节码
     * @param transformation 算法变种
     * @param ivLength 初始向量长度.此值视不同的区块加解密模式而定,除默认SM4算法变种和ECB模式为0外,
     *            其他的模式由长度为 12、15、16位不等,大部分为16位
     * @return 原文字节码
     * @throws RuntimeException */
    public static byte[] decrypt(byte[] keyBytes, byte[] cipherBytes, String transformation) {

        int index = of(transformation).get().indexOf('/');
        int ivLength = 0;

        // 排除 默认SM4算法变种和 ECB 工作模式没有 IV外，其他的 IV 长度等于密文的第一个字节
        if ((index > 0)
                && !(Mode.ECB.toString().equals(transformation.substring(index + 1, transformation.lastIndexOf('/')))))
            ivLength = of(cipherBytes).get()[0];

        try {
            Cipher cipher = Cipher.getInstance(transformation, PROVIDER_NAME);

            if (ivLength > 0) {
                // logger.info("接收方收到密文:IV+Cipher({} bytes): {}", cipherBytes.length, Encoder.Hex.encode(cipherBytes));
                byte[] ivBytes = new byte[ivLength];
                System.arraycopy(cipherBytes, 1, ivBytes, 0, ivLength);
                byte[] cipherBytesNoIvLenghtAndIv = new byte[cipherBytes.length - ivLength - 1];
                System.arraycopy(cipherBytes, ivLength + 1, cipherBytesNoIvLenghtAndIv, 0,
                        cipherBytes.length - ivLength - 1);
                cipherBytes = cipherBytesNoIvLenghtAndIv;
                // logger.info("接收方提取初始向量:(IV {} bytes): {}", ivBytes.length, Encoder.Hex.encode(ivBytes));
                cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(of(keyBytes).get(), transformation),
                        new IvParameterSpec(ivBytes));
            } else {
                cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(of(keyBytes).get(), transformation));
            }

            return cipher.doFinal(cipherBytes);
        } catch (GeneralSecurityException e) {
            logger.warn("SM4 Decrypt Fil.Reason: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }
}
