package weixin_ilink

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

// ErrBadCiphertext is returned when ciphertext is not a whole number of AES
// blocks or its padding is malformed.
var ErrBadCiphertext = errors.New("weixin_ilink: bad ciphertext or padding")

// pkcs7Pad appends PKCS#7 padding to src so its length becomes a multiple of
// blockSize. Every appended byte holds the pad length, so it self-describes
// and Unpad can strip it unambiguously.
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 strips PKCS#7 padding from src. It returns ErrBadCiphertext when
// the length is not a block multiple or the final pad byte is inconsistent.
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 encrypts src under AES with the supplied key using ECB mode.
// Go's crypto/cipher package does not ship ECB, so the mode is implemented
// here: src is padded to a block multiple and every block is encrypted
// independently with the same key. ECB is deterministic and not semantically
// secure; it is required here only because the iLink CDN protocol mandates it
// for media payloads.
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 decrypts src under AES with the supplied key and strips the
// PKCS#7 padding. ECB has no IV, so the first block decrypts directly with the
// key. It returns ErrBadCiphertext for malformed input.
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())
}