package com.qmth.boot.tools.crypto;

import com.qmth.boot.tools.models.ByteArray;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.util.WeakHashMap;

/**
 * AES加解密算法工具，使用默认的PKCS5Padding填充
 */
public class AES {

    private static final String ALGORITHM = "AES";

    private static final String ECB_ALGORITHM = "AES/ECB/PKCS5Padding";

    private static final String CBC_ALGORITHM = "AES/CBC/PKCS5Padding";

    private static WeakHashMap<String, SecretKeySpec> keyMap = new WeakHashMap<>();

    private static WeakHashMap<String, IvParameterSpec> vectorMap = new WeakHashMap<>();

    /**
     * 获取ECB模式加密工具
     *
     * @param key
     * @return
     * @throws Exception
     */
    public static Cipher getEncryptCipher(String key) throws Exception {
        Cipher cipher = Cipher.getInstance(ECB_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, getSecretKey(key));
        return cipher;
    }

    /**
     * 获取ECB模式解密工具
     *
     * @param key
     * @return
     * @throws Exception
     */
    public static Cipher getDecryptCipher(String key) throws Exception {
        Cipher cipher = Cipher.getInstance(ECB_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, getSecretKey(key));
        return cipher;
    }

    /**
     * ECB模式加密
     *
     * @param data
     * @param key
     * @return
     */
    public static ByteArray encrypt(byte[] data, String key) throws Exception {
        Cipher cipher = Cipher.getInstance(ECB_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, getSecretKey(key));
        return ByteArray.fromArray(cipher.doFinal(data));
    }

    /**
     * ECB模式解密
     *
     * @param data
     * @param key
     * @return
     */
    public static ByteArray decrypt(byte[] data, String key) throws Exception {
        Cipher cipher = Cipher.getInstance(ECB_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, getSecretKey(key));
        return ByteArray.fromArray(cipher.doFinal(data));
    }

    /**
     * CBC模式加密
     *
     * @param data
     * @param key
     * @param vector
     * @return
     */
    public static ByteArray encrypt(byte[] data, String key, String vector) throws Exception {
        Cipher cipher = Cipher.getInstance(CBC_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, getSecretKey(key), getVector(vector));
        return ByteArray.fromArray(cipher.doFinal(data));
    }

    /**
     * CBC模式解密
     *
     * @param data
     * @param key
     * @param vector
     * @return
     */
    public static ByteArray decrypt(byte[] data, String key, String vector) throws Exception {
        Cipher cipher = Cipher.getInstance(CBC_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, getSecretKey(key), getVector(vector));
        return ByteArray.fromArray(cipher.doFinal(data));
    }

    private static SecretKeySpec getSecretKey(String key) {
        return keyMap.computeIfAbsent(key, k -> new SecretKeySpec(k.getBytes(StandardCharsets.UTF_8), ALGORITHM));
    }

    private static IvParameterSpec getVector(String vector) {
        return vectorMap.computeIfAbsent(vector, k -> new IvParameterSpec(k.getBytes(StandardCharsets.UTF_8)));
    }
}
