package com.bjlx.armcloudaf;

import java.io.ByteArrayOutputStream;
import java.util.Arrays;

import javax.crypto.Cipher;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

public class AESPatchHelper {
    private static final byte[] SALT = new byte[8]; // 全零 salt
    private static final int PADDING_LEN = 8;
    private static final int IV_LEN = 16;

    public static byte[] extractIV(byte[] data) {
        int ivPos = data.length - PADDING_LEN - IV_LEN;
        return Arrays.copyOfRange(data, ivPos, ivPos + IV_LEN);
    }

    public static byte[] extractPadding(byte[] data) {
        int padPos = data.length - PADDING_LEN - IV_LEN - PADDING_LEN;
        if (padPos < 0) padPos = data.length - PADDING_LEN - IV_LEN;
        return Arrays.copyOfRange(data, data.length - PADDING_LEN - IV_LEN, data.length - IV_LEN);
    }

    public static byte[] extractCiphertext(byte[] data) {
        int ctLen = data.length - PADDING_LEN - IV_LEN;
        return Arrays.copyOfRange(data, 0, ctLen);
    }

    private static SecretKeySpec deriveKey(String password) throws Exception {
        return new SecretKeySpec(
                SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1")
                        .generateSecret(new PBEKeySpec(password.toCharArray(), SALT, 10000, 128))
                        .getEncoded(),
                "AES"
        );
    }

    public static byte[] decryptRaw(byte[] ciphertext, byte[] iv, String password) throws Exception {
        SecretKeySpec key = deriveKey(password);
        Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
        cipher.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));
        return cipher.doFinal(ciphertext);
    }

    public static byte[] encryptRaw(byte[] plaintext, byte[] iv, String password) throws Exception {
        SecretKeySpec key = deriveKey(password);
        Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
        cipher.init(Cipher.ENCRYPT_MODE, key, new IvParameterSpec(iv));
        return cipher.doFinal(plaintext);
    }

    /**
     * 入口函数：
     * 提取 data 中 IV + padding + 密文 → 解密 → 加密 → 再拼装成结构化数据
     */
    public static byte[] patchAndReEncrypt(byte[] data, String password) throws Exception {
        if (data == null || data.length <= (PADDING_LEN + IV_LEN)) return null;

        byte[] iv = extractIV(data);
        byte[] padding = extractPadding(data);
        byte[] ciphertext = extractCiphertext(data);
        // 解密
        byte[] plaintext = decryptRaw(ciphertext, iv, password);
        // 可在此修改 plaintext（如篡改某些字段）
        // 再加密
        byte[] newCiphertext = encryptRaw(plaintext, iv, password);

        // 拼装为结构：[newCiphertext] + [padding] + [iv]
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        out.write(newCiphertext);
        out.write(padding);
        out.write(iv);
        return out.toByteArray();
    }
}