// AES加解密：请求体加密
package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"simple-memo/global"
)

// 注意：这里不再直接初始化，等到配置加载后再赋值
var AESKey []byte
var AESIV []byte

// InitAES 初始化AES密钥（必须在配置加载后调用）
func InitAES() {
	AESKey = []byte(global.VP.GetString("aes.key")) // 32位
	AESIV = []byte(global.VP.GetString("aes.iv"))   // 16位
}

// PKCS7Padding 填充（修复完成）
func PKCS7Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padText...) // 这里加 ... 即可修复
}

// AESEncrypt 加密
func AESEncrypt(origData []byte) (string, error) {
	block, err := aes.NewCipher(AESKey)
	if err != nil {
		return "", err
	}
	blockSize := block.BlockSize()
	origData = PKCS7Padding(origData, blockSize)
	blockMode := cipher.NewCBCEncrypter(block, AESIV)
	encrypted := make([]byte, len(origData))
	blockMode.CryptBlocks(encrypted, origData)
	return base64.StdEncoding.EncodeToString(encrypted), nil
}

// AESDecrypt 解密
func AESDecrypt(encrypted string) ([]byte, error) {
	block, err := aes.NewCipher(AESKey)
	if err != nil {
		return nil, err
	}
	decoded, err := base64.StdEncoding.DecodeString(encrypted)
	if err != nil {
		return nil, err
	}
	blockMode := cipher.NewCBCDecrypter(block, AESIV)
	orig := make([]byte, len(decoded))
	blockMode.CryptBlocks(orig, decoded)

	// ========= 这里换成安全版去填充 =========
	orig = PKCS7UnPaddingSafe(orig)
	return orig, nil
}

// PKCS7UnPadding 去填充（原版保留，不删）
func PKCS7UnPadding(origData []byte) []byte {
	length := len(origData)
	if length == 0 {
		return []byte{} // 空数据直接返回，不崩溃
	}
	unPadding := int(origData[length-1])
	return origData[:length-unPadding]
}

// PKCS7UnPaddingSafe 安全去填充，永不 panic
func PKCS7UnPaddingSafe(origData []byte) []byte {
	length := len(origData)
	if length == 0 {
		return []byte{}
	}

	unPadding := int(origData[length-1])

	// ========= 安全校验：防止 slice 越界崩溃 =========
	if unPadding < 1 || unPadding > length {
		return origData
	}

	return origData[:length-unPadding]
}