package cn.gintone.encryptionUtils;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import java.io.ByteArrayOutputStream;
import java.io.ObjectOutputStream;
import java.nio.charset.StandardCharsets;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.util.Base64;

public class SecureHybridEncryptor {
    private static final int AES_KEY_SIZE = 128;
    private static final int GCM_TAG_LENGTH = 128;
    private static final int GCM_IV_LENGTH = 12; // bytes

    // 主加密方法
    public static String encrypt(String plaintext, PublicKey publicKey) throws Exception {
        // 1. 生成随机AES密钥
        SecretKey aesKey = generateAESKey();

        // 2. 使用AES加密原始数据
        byte[] iv = generateSecureIV();
        String[] aesResult = encryptWithAES(plaintext, aesKey, iv);
        String encryptedData = aesResult[0];
        String ivBase64 = aesResult[1];

        // 3. 使用RSA加密AES密钥
        String encryptedAESKey = encryptAESKey(aesKey, publicKey);

        // 4. 组合加密结果
        return formatFinalResult(encryptedAESKey, ivBase64, encryptedData);
    }

    public static String encrypt_object(Object obj, PublicKey publicKey) throws Exception {
        // 1. 生成随机AES密钥
        SecretKey aesKey = generateAESKey();

        // 2. 使用AES加密原始数据
        byte[] iv = generateSecureIV();
        String[] aesResult = encryptWithAES_object(obj, aesKey, iv);
        String encryptedData = aesResult[0];
        String ivBase64 = aesResult[1];

        // 3. 使用RSA加密AES密钥
        String encryptedAESKey = encryptAESKey(aesKey, publicKey);

        // 4. 组合加密结果
        return formatFinalResult(encryptedAESKey, ivBase64, encryptedData);
    }

    private static SecretKey generateAESKey() throws Exception {
        KeyGenerator keyGen = KeyGenerator.getInstance("AES");
        keyGen.init(AES_KEY_SIZE);
        return keyGen.generateKey();
    }

    private static byte[] generateSecureIV() {
        byte[] iv = new byte[GCM_IV_LENGTH];
        new SecureRandom().nextBytes(iv);
        return iv;
    }


    private static String[] encryptWithAES(String plaintext, SecretKey aesKey, byte[] iv) throws Exception {
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
        cipher.init(Cipher.ENCRYPT_MODE, aesKey, spec);

        byte[] encrypted = cipher.doFinal(plaintext.getBytes(StandardCharsets.UTF_8));
        return new String[] {
                Base64.getEncoder().encodeToString(encrypted),
                Base64.getEncoder().encodeToString(iv)
        };
    }

    private static String[] encryptWithAES_object(Object obj, SecretKey aesKey, byte[] iv) throws Exception {

        // 将对象序列化为字节数组
        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
        ObjectOutputStream objectOutputStream = new ObjectOutputStream(byteArrayOutputStream);
        objectOutputStream.writeObject(obj);
        objectOutputStream.flush();
        byte[] serializedObject = byteArrayOutputStream.toByteArray();
        objectOutputStream.close();


        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
        cipher.init(Cipher.ENCRYPT_MODE, aesKey, spec);

        byte[] encrypted = cipher.doFinal(serializedObject);
        return new String[] {
                Base64.getEncoder().encodeToString(encrypted),
                Base64.getEncoder().encodeToString(iv)
        };
    }

    private static String encryptAESKey(SecretKey aesKey, PublicKey publicKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] encryptedKey = cipher.doFinal(aesKey.getEncoded());
        return Base64.getEncoder().encodeToString(encryptedKey);
    }

    private static String formatFinalResult(String encryptedKey, String iv, String data) {
        return encryptedKey + ":" + iv + ":" + data;
    }
}