package com.chinaweal.aiccs.common.util;

import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.modes.CBCBlockCipher;
import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.util.encoders.Hex;

import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;

/**
 *
 * @author lroyia
 * @since 2026/1/4 15:33
 **/
public class SM4Utils {


    private static final int BLOCK_SIZE = 16;

    /**
     * SM4 加密
     *
     * @param plaintext 明文字符串
     * @param sm4Key SM4密钥
     * @return 十六进制格式的密文字符串
     */
    public static String encrypt(String plaintext, String sm4Key) {
        try {
            byte[] keyBytes = Hex.decode(sm4Key);
            byte[] plaintextBytes = plaintext.getBytes(StandardCharsets.UTF_8);

            // 生成随机 IV
            byte[] iv = new byte[BLOCK_SIZE];
            SecureRandom random = new SecureRandom();
            random.nextBytes(iv);

            // 执行加密
            SM4Engine engine = new SM4Engine();
            CBCBlockCipher blockCipher = new CBCBlockCipher(engine);
            PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(blockCipher);

            KeyParameter keyParameter = new KeyParameter(keyBytes);
            ParametersWithIV params = new ParametersWithIV(keyParameter, iv);
            cipher.init(true, params);

            byte[] encrypted = new byte[cipher.getOutputSize(plaintextBytes.length)];
            int len = cipher.processBytes(plaintextBytes, 0, plaintextBytes.length, encrypted, 0);
            len += cipher.doFinal(encrypted, len);

            // 将 IV 和密文合并返回
            byte[] result = new byte[iv.length + len];
            System.arraycopy(iv, 0, result, 0, iv.length);
            System.arraycopy(encrypted, 0, result, iv.length, len);

            return Hex.toHexString(result);
        } catch (Exception e) {
            throw new RuntimeException("SM4 加密失败", e);
        }
    }

    /**
     * SM4 解密
     *
     * @param ciphertext 十六进制格式的密文字符串
     * @return 解密后的明文字符串
     */
    public static String decrypt(String ciphertext, String sm4Key) {
        try {
            byte[] keyBytes = Hex.decode(sm4Key);
            byte[] ciphertextBytes = Hex.decode(ciphertext);

            if (ciphertextBytes.length < BLOCK_SIZE) {
                throw new IllegalArgumentException("密文长度不足");
            }

            // 提取 IV 和密文
            byte[] iv = new byte[BLOCK_SIZE];
            byte[] encrypted = new byte[ciphertextBytes.length - BLOCK_SIZE];
            System.arraycopy(ciphertextBytes, 0, iv, 0, BLOCK_SIZE);
            System.arraycopy(ciphertextBytes, BLOCK_SIZE, encrypted, 0, encrypted.length);

            // 执行解密
            SM4Engine engine = new SM4Engine();
            CBCBlockCipher blockCipher = new CBCBlockCipher(engine);
            PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(blockCipher);

            KeyParameter keyParameter = new KeyParameter(keyBytes);
            ParametersWithIV params = new ParametersWithIV(keyParameter, iv);
            cipher.init(false, params);

            byte[] decrypted = new byte[cipher.getOutputSize(encrypted.length)];
            int len = cipher.processBytes(encrypted, 0, encrypted.length, decrypted, 0);
            len += cipher.doFinal(decrypted, len);

            // 移除填充并返回明文
            byte[] result = new byte[len];
            System.arraycopy(decrypted, 0, result, 0, len);

            return new String(result, StandardCharsets.UTF_8);
        } catch (Exception e) {
            throw new RuntimeException("SM4 解密失败", e);
        }
    }
}
