package weixin_ilink

import (
	"crypto/aes"
	"errors"
	"fmt"
)

// ErrBadCiphertext 在密文长度不是 AES 分组整数倍、或其填充格式非法时返回。
var ErrBadCiphertext = errors.New("weixin_ilink: bad ciphertext or padding")

// pkcs7Pad 按 PKCS#7 向 src 追加填充，使其长度成为 blockSize 的整数倍。
// 每个追加字节都写入填充长度，因此它是自描述的，Unpad 可以无歧义地剥离。
func pkcs7Pad(src []byte, blockSize int) []byte {
	n := blockSize - len(src)%blockSize
	out := make([]byte, len(src)+n)
	copy(out, src)
	for i := len(src); i < len(out); i++ {
		out[i] = byte(n)
	}
	return out
}

// pkcs7Unpad 从 src 剥离 PKCS#7 填充。长度不是分组整数倍、或末尾填充字节不自洽
// 时返回 ErrBadCiphertext。
func pkcs7Unpad(src []byte, blockSize int) ([]byte, error) {
	if len(src) == 0 || len(src)%blockSize != 0 {
		return nil, ErrBadCiphertext
	}
	n := int(src[len(src)-1])
	if n < 1 || n > blockSize {
		return nil, ErrBadCiphertext
	}
	for _, b := range src[len(src)-n:] {
		if int(b) != n {
			return nil, ErrBadCiphertext
		}
	}
	return src[:len(src)-n], nil
}

// aesECBEncrypt 用给定密钥以 ECB 模式对 src 做 AES 加密。Go 的 crypto/cipher
// 不提供 ECB，所以此模式在此自行实现：先把 src 填充到分组整数倍，再用同一密钥
// 逐块独立加密。ECB 是确定性的、不具备语义安全性；这里之所以需要它，仅仅是因为
// iLink CDN 协议对媒体负载强制要求该模式。
func aesECBEncrypt(key, src []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("weixin_ilink: new aes cipher: %w", err)
	}
	padded := pkcs7Pad(src, block.BlockSize())
	out := make([]byte, len(padded))
	for i := 0; i < len(padded); i += block.BlockSize() {
		block.Encrypt(out[i:i+block.BlockSize()], padded[i:i+block.BlockSize()])
	}
	return out, nil
}

// aesECBDecrypt 用给定密钥对 src 做 AES 解密并剥离 PKCS#7 填充。ECB 没有 IV，
// 所以首块直接用密钥解密。输入格式非法时返回 ErrBadCiphertext。
func aesECBDecrypt(key, src []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("weixin_ilink: new aes cipher: %w", err)
	}
	if len(src) == 0 || len(src)%block.BlockSize() != 0 {
		return nil, ErrBadCiphertext
	}
	out := make([]byte, len(src))
	for i := 0; i < len(src); i += block.BlockSize() {
		block.Decrypt(out[i:i+block.BlockSize()], src[i:i+block.BlockSize()])
	}
	return pkcs7Unpad(out, block.BlockSize())
}